diff --git a/api/commands/retention.py b/api/commands/retention.py index 47eaa76c957..d03c9bcc6da 100644 --- a/api/commands/retention.py +++ b/api/commands/retention.py @@ -1,6 +1,8 @@ import datetime import logging +import re import time +import uuid from collections.abc import Callable from typing import TypedDict @@ -21,6 +23,7 @@ from tasks.remove_app_and_related_data_task import delete_draft_variables_batch logger = logging.getLogger(__name__) _HEX_PREFIXES = tuple("0123456789abcdef") +_TARGET_MONTH_PATTERN = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$") class WorkflowRunArchivePlanRow(TypedDict): @@ -66,6 +69,7 @@ def _parse_tenant_prefixes(prefixes: str | None) -> list[str]: def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[str] | None: + """Keep an omitted scope unset while rejecting an explicitly empty scope.""" if raw_ids is None: return None parsed = sorted({raw_id.strip() for raw_id in raw_ids.split(",") if raw_id.strip()}) @@ -74,6 +78,27 @@ def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[ return parsed +def _parse_archive_target_month(target_month: str) -> tuple[int, int]: + """Validate the V2 catalog month selector and return its numeric components.""" + if not _TARGET_MONTH_PATTERN.fullmatch(target_month): + raise click.BadParameter("target-month must use YYYY-MM format", param_hint="--target-month") + year_text, month_text = target_month.split("-", maxsplit=1) + return int(year_text), int(month_text) + + +def _parse_archive_catalog_cursor(after_catalog_id: str | None) -> str | None: + """Normalize the exclusive V2 catalog keyset cursor when one is provided.""" + if after_catalog_id is None: + return None + try: + return str(uuid.UUID(after_catalog_id)) + except ValueError as exc: + raise click.BadParameter( + "after-catalog-id must be a UUID returned by the same V2 operation and scope", + param_hint="--after-catalog-id", + ) from exc + + def _get_archive_candidate_tenant_ids_by_prefix( session: Session, prefix: str, @@ -810,9 +835,11 @@ def backfill_workflow_run_archive_bundles( click.echo(click.style(f" ... and {len(summary.errors) - 10} more failures", fg="red")) -def _echo_bundle_archive_operation_summary(summary) -> None: +def _echo_bundle_archive_operation_summary(summary, *, dry_run: bool) -> None: status = "completed successfully" if summary.bundles_failed == 0 else "completed with failures" fg = "green" if summary.bundles_failed == 0 else "red" + cursor_label = "preview_next_catalog_id" if dry_run else "next_catalog_id" + cursor_value = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id click.echo( click.style( f"{summary.operation} {status}. " @@ -821,10 +848,12 @@ def _echo_bundle_archive_operation_summary(summary) -> None: f"archive_bytes={summary.archive_bytes} duration={summary.elapsed_time:.2f}s " f"validation_time={summary.validation_time:.2f}s " f"runs_per_second={summary.runs_per_second:.2f} rows_per_second={summary.rows_per_second:.2f} " - f"bytes_per_second={summary.bytes_per_second:.2f}", + f"bytes_per_second={summary.bytes_per_second:.2f} {cursor_label}={cursor_value or 'none'}", fg=fg, ) ) + if dry_run: + click.echo(click.style("Dry-run cursor is preview-only; do not persist it for a destructive run.", fg="yellow")) click.echo(click.style("table,row_count", fg="white")) for table_name in [ "workflow_runs", @@ -842,7 +871,8 @@ def _echo_bundle_archive_operation_summary(summary) -> None: click.style( f" bundle={result.bundle_id} tenant={result.tenant_id} runs={result.run_count} " f"rows={result.row_count} archive_bytes={result.archive_bytes} " - f"time={result.elapsed_time:.2f}s validation={result.validation_time:.2f}s", + f"catalog_id={result.catalog_id} time={result.elapsed_time:.2f}s " + f"validation={result.validation_time:.2f}s", fg="white", ) ) @@ -850,7 +880,7 @@ def _echo_bundle_archive_operation_summary(summary) -> None: click.echo( click.style( f" failed bundle={result.bundle_id} tenant={result.tenant_id} " - f"object_prefix={result.object_prefix} error={result.error}", + f"catalog_id={result.catalog_id} object_prefix={result.object_prefix} error={result.error}", fg="red", ) ) @@ -867,25 +897,24 @@ def _echo_bundle_archive_operation_summary(summary) -> None: ) @click.option("--run-id", required=False, help="Workflow run ID to restore.") @click.option( - "--start-from", - type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + "--target-month", + metavar="YYYY-MM", default=None, - help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.", + help="V2 catalog month to restore; required unless --run-id is used.", ) @click.option( - "--end-before", - type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + "--after-catalog-id", default=None, - help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.", + help="Exclusive V2 cursor from the same restore month and tenant scope.", ) @click.option("--workers", default=1, show_default=True, type=int, help="V1 --run-id compatibility only.") -@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to restore.") +@click.option("--limit", type=click.IntRange(min=1), default=100, show_default=True, help="Maximum V2 catalog rows.") @click.option("--dry-run", is_flag=True, help="Preview without restoring.") def restore_workflow_runs( tenant_ids: str | None, run_id: str | None, - start_from: datetime.datetime | None, - end_before: datetime.datetime | None, + target_month: str | None, + after_catalog_id: str | None, workers: int, limit: int, dry_run: bool, @@ -905,23 +934,20 @@ def restore_workflow_runs( from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance from services.retention.workflow_run.restore_archived_workflow_run import WorkflowRunRestore - parsed_tenant_ids = None - if tenant_ids: - parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()] - if not parsed_tenant_ids: - raise click.BadParameter("tenant-ids must not be empty") + parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids") - if (start_from is None) ^ (end_before is None): - raise click.UsageError("--start-from and --end-before must be provided together.") - if run_id is None and (start_from is None or end_before is None): - raise click.UsageError("--start-from and --end-before are required for batch restore.") if workers < 1: raise click.BadParameter("workers must be at least 1") + if run_id is not None and (target_month is not None or after_catalog_id is not None): + raise click.UsageError("--target-month and --after-catalog-id are only valid for V2 batch restore.") + if run_id is None and target_month is None: + raise click.UsageError("--target-month is required for V2 batch restore.") start_time = datetime.datetime.now(datetime.UTC) + target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}" click.echo( click.style( - f"Starting restore of workflow run {run_id} at {start_time.isoformat()}.", + f"Starting restore of {target_desc} at {start_time.isoformat()}.", fg="white", ) ) @@ -955,17 +981,20 @@ def restore_workflow_runs( click.echo( click.style("--workers is ignored for V2 bundle restore; bundles are processed serially.", fg="yellow") ) - assert start_from is not None - assert end_before is not None + assert target_month is not None + target_year, target_month_number = _parse_archive_target_month(target_month) + catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id) bundle_restorer = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True) summary = bundle_restorer.restore_batch( tenant_ids=parsed_tenant_ids, - start_date=start_from, - end_date=end_before, + target_year=target_year, + target_month=target_month_number, + after_catalog_id=catalog_cursor, limit=limit, ) - _echo_bundle_archive_operation_summary(summary) - return + _echo_bundle_archive_operation_summary(summary, dry_run=dry_run) + if summary.bundles_failed: + raise click.exceptions.Exit(1) @click.command( @@ -979,23 +1008,41 @@ def restore_workflow_runs( ) @click.option("--run-id", required=False, help="Workflow run ID to delete.") @click.option( - "--start-from", - type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + "--target-month", + metavar="YYYY-MM", default=None, - help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.", + help="V2 catalog month to delete; required unless --run-id is used.", ) @click.option( - "--end-before", - type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]), + "--after-catalog-id", default=None, - help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.", + help="Exclusive V2 cursor from the same delete month and tenant scope.", +) +@click.option( + "--run-shard-index", + default=None, + type=click.IntRange(min=0), + help="Zero-based archive shard index. Must be paired with --run-shard-total.", +) +@click.option( + "--run-shard-total", + default=None, + type=click.IntRange(min=1, max=16), + help="Total archive shard count. Must be paired with --run-shard-index.", +) +@click.option("--all-pages", is_flag=True, help="Process catalog pages until an empty page is reached.") +@click.option( + "--limit", + type=click.IntRange(min=1), + default=100, + show_default=True, + help="Maximum V2 catalog rows per page.", ) -@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to delete.") @click.option("--dry-run", is_flag=True, help="Preview without deleting.") @click.option( "--skip-bad-archives", is_flag=True, - help="Continue batch deletion when one archive object fails validation.", + help="V1 --run-id only: continue when one archive object fails validation.", ) @click.option( "--restore-sample-interval", @@ -1007,8 +1054,11 @@ def restore_workflow_runs( def delete_archived_workflow_runs( tenant_ids: str | None, run_id: str | None, - start_from: datetime.datetime | None, - end_before: datetime.datetime | None, + target_month: str | None, + after_catalog_id: str | None, + run_shard_index: int | None, + run_shard_total: int | None, + all_pages: bool, limit: int, dry_run: bool, skip_bad_archives: bool, @@ -1018,26 +1068,38 @@ def delete_archived_workflow_runs( Delete archived workflow runs from the database. Batch delete uses V2 bundle metadata and validates object existence, manifest schema, object size, checksum, row - counts, and source/archive content checksums before deleting source rows. `--run-id` keeps the V1 per-run path. + counts, and source/archive content checksums before deleting source rows. Parallel workers may select one exact + archive shard; all-pages mode keeps only the current bounded page in memory. `--run-id` keeps the V1 per-run path. """ from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion - parsed_tenant_ids = None - if tenant_ids: - parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()] - if not parsed_tenant_ids: - raise click.BadParameter("tenant-ids must not be empty") + parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids") - if (start_from is None) ^ (end_before is None): - raise click.UsageError("--start-from and --end-before must be provided together.") - if run_id is None and (start_from is None or end_before is None): - raise click.UsageError("--start-from and --end-before are required for batch delete.") if restore_sample_interval < 0: raise click.BadParameter("restore-sample-interval must be >= 0") + if run_id is not None and ( + target_month is not None + or after_catalog_id is not None + or run_shard_index is not None + or run_shard_total is not None + or all_pages + ): + raise click.UsageError( + "--target-month, --after-catalog-id, --run-shard-index, --run-shard-total, and --all-pages " + "are only valid for V2 batch delete." + ) + if run_id is None and target_month is None: + raise click.UsageError("--target-month is required for V2 batch delete.") + if run_id is None and skip_bad_archives: + raise click.UsageError("--skip-bad-archives is not supported for V2 catalog batches; they fail fast.") + if (run_shard_index is None) ^ (run_shard_total is None): + raise click.UsageError("--run-shard-index and --run-shard-total must be provided together.") + if run_shard_index is not None and run_shard_total is not None and run_shard_index >= run_shard_total: + raise click.UsageError("--run-shard-index must be less than --run-shard-total.") start_time = datetime.datetime.now(datetime.UTC) - target_desc = f"workflow run {run_id}" if run_id else "workflow runs" + target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}" click.echo( click.style( f"Starting delete of {target_desc} at {start_time.isoformat()}.", @@ -1110,20 +1172,104 @@ def delete_archived_workflow_runs( if restore_sample_interval: click.echo(click.style("--restore-sample-interval is ignored for V2 bundle delete.", fg="yellow")) - assert start_from is not None - assert end_before is not None - bundle_deleter = WorkflowRunBundleArchiveMaintenance( - dry_run=dry_run, - strict_content_validation=True, - stop_on_error=not skip_bad_archives, + assert target_month is not None + target_year, target_month_number = _parse_archive_target_month(target_month) + catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id) + shard = ( + f"{run_shard_index:02d}-of-{run_shard_total:02d}" + if run_shard_index is not None and run_shard_total is not None + else None ) - summary = bundle_deleter.delete_batch( - tenant_ids=parsed_tenant_ids, - start_date=start_from, - end_date=end_before, - limit=limit, - ) - _echo_bundle_archive_operation_summary(summary) + bundle_deleter = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True) + if run_shard_total is not None: + try: + bundle_deleter.validate_catalog_shards( + target_year=target_year, + target_month=target_month_number, + shard_total=run_shard_total, + tenant_ids=parsed_tenant_ids, + ) + except ValueError as exc: + logger.exception( + "Archive catalog shard preflight failed: target_month=%s shard=%s", + target_month, + shard, + ) + raise click.ClickException( + f"Archive catalog shard preflight failed for target_month={target_month} shard={shard}: {exc}" + ) from exc + + initial_catalog_cursor = catalog_cursor + pages_processed = 0 + bundles_succeeded = 0 + runs_processed = 0 + rows_processed = 0 + archive_bytes = 0 + while True: + summary = bundle_deleter.delete_batch( + tenant_ids=parsed_tenant_ids, + target_year=target_year, + target_month=target_month_number, + after_catalog_id=catalog_cursor, + limit=limit, + shard=shard, + ) + _echo_bundle_archive_operation_summary(summary, dry_run=dry_run) + if summary.bundles_failed: + failed_result = next((result for result in summary.results if not result.success), None) + failed_catalog_id = failed_result.catalog_id if failed_result is not None else "unknown" + page_resume_cursor = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id + resume_cursor = page_resume_cursor or catalog_cursor + if dry_run: + cursor_details = ( + f"preview_after_catalog_id={resume_cursor or 'none'} " + f"destructive_retry_after_catalog_id={initial_catalog_cursor or 'none'}" + ) + else: + cursor_details = f"resume_after_catalog_id={resume_cursor or 'none'}" + click.echo( + click.style( + f"Delete stopped: target_month={target_month} shard={shard or 'all'} " + f"failed_catalog_id={failed_catalog_id} " + f"{cursor_details}", + fg="red", + ) + ) + raise click.exceptions.Exit(1) + + if not all_pages: + break + if summary.bundles_processed == 0: + break + + pages_processed += 1 + bundles_succeeded += summary.bundles_succeeded + runs_processed += summary.runs_processed + rows_processed += summary.rows_processed + archive_bytes += summary.archive_bytes + next_catalog_id = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id + if next_catalog_id is None or (catalog_cursor is not None and next_catalog_id <= catalog_cursor): + click.echo( + click.style( + f"Delete cursor did not advance: target_month={target_month} shard={shard or 'all'} " + f"after_catalog_id={catalog_cursor or 'none'} next_catalog_id={next_catalog_id or 'none'}", + fg="red", + ) + ) + raise click.exceptions.Exit(1) + catalog_cursor = next_catalog_id + + if all_pages: + final_cursor_label = "preview_final_catalog_id" if dry_run else "final_catalog_id" + click.echo( + click.style( + f"Delete all-pages completed successfully. target_month={target_month} shard={shard or 'all'} " + f"pages={pages_processed} bundles_success={bundles_succeeded} runs={runs_processed} " + f"rows={rows_processed} archive_bytes={archive_bytes} " + f"{final_cursor_label}={catalog_cursor or 'none'}", + fg="green", + ) + ) def _find_orphaned_draft_variables(batch_size: int = 1000) -> list[str]: diff --git a/api/libs/archive_storage.py b/api/libs/archive_storage.py index fb35a52948a..5df13830b35 100644 --- a/api/libs/archive_storage.py +++ b/api/libs/archive_storage.py @@ -16,12 +16,19 @@ from urllib.parse import quote import boto3 import orjson from botocore.client import Config -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from configs import dify_config logger = logging.getLogger(__name__) +_OBJECT_NOT_FOUND_ERROR_CODES = frozenset({"404", "NoSuchKey", "NotFound"}) + + +def _is_object_not_found_error(error: ClientError) -> bool: + error_code = str(error.response.get("Error", {}).get("Code", "")) + return error_code in _OBJECT_NOT_FOUND_ERROR_CODES + class ArchiveStorageError(Exception): """Base exception for archive storage operations.""" @@ -138,10 +145,11 @@ class ArchiveStorage: response = self.client.get_object(Bucket=self.bucket, Key=key) return response["Body"].read() except ClientError as e: - error_code = e.response.get("Error", {}).get("Code") - if error_code == "NoSuchKey": - raise FileNotFoundError(f"Archive object not found: {key}") - raise ArchiveStorageError(f"Failed to download object '{key}': {e}") + if _is_object_not_found_error(e): + raise FileNotFoundError(f"Archive object not found: {key}") from e + raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e + except BotoCoreError as e: + raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e def get_object_stream(self, key: str) -> Generator[bytes, None, None]: """ @@ -161,10 +169,11 @@ class ArchiveStorage: response = self.client.get_object(Bucket=self.bucket, Key=key) yield from response["Body"].iter_chunks() except ClientError as e: - error_code = e.response.get("Error", {}).get("Code") - if error_code == "NoSuchKey": - raise FileNotFoundError(f"Archive object not found: {key}") - raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") + if _is_object_not_found_error(e): + raise FileNotFoundError(f"Archive object not found: {key}") from e + raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e + except BotoCoreError as e: + raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e def object_exists(self, key: str) -> bool: """ @@ -175,12 +184,19 @@ class ArchiveStorage: Returns: True if object exists, False otherwise + + Raises: + ArchiveStorageError: If storage cannot authoritatively determine object existence """ try: self.client.head_object(Bucket=self.bucket, Key=key) return True - except ClientError: - return False + except ClientError as e: + if _is_object_not_found_error(e): + return False + raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e + except BotoCoreError as e: + raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e def delete_object(self, key: str) -> None: """ @@ -196,7 +212,12 @@ class ArchiveStorage: self.client.delete_object(Bucket=self.bucket, Key=key) logger.debug("Deleted object: %s", key) except ClientError as e: - raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") + if _is_object_not_found_error(e): + logger.debug("Archive object was already absent: %s", key) + return + raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e + except BotoCoreError as e: + raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e def generate_presigned_url( self, diff --git a/api/migrations/versions/2026_07_15_1600-3c9f8e2a1d7b_add_workflow_run_archive_bundle_cursor_indexes.py b/api/migrations/versions/2026_07_15_1600-3c9f8e2a1d7b_add_workflow_run_archive_bundle_cursor_indexes.py new file mode 100644 index 00000000000..8660c19aa14 --- /dev/null +++ b/api/migrations/versions/2026_07_15_1600-3c9f8e2a1d7b_add_workflow_run_archive_bundle_cursor_indexes.py @@ -0,0 +1,45 @@ +"""add workflow-run archive bundle cursor indexes + +Revision ID: 3c9f8e2a1d7b +Revises: 7a1c2d9e4b60 +Create Date: 2026-07-15 16:00:00.000000 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "3c9f8e2a1d7b" +down_revision = "7a1c2d9e4b60" +branch_labels = None +depends_on = None + +_TABLE_NAME = "workflow_run_archive_bundles" +_INDEXES = ( + ("workflow_run_archive_bundle_month_id_idx", ("year", "month", "id")), + ("workflow_run_archive_bundle_month_shard_id_idx", ("year", "month", "shard", "id")), +) + + +def _is_postgresql() -> bool: + return op.get_bind().dialect.name == "postgresql" + + +def upgrade() -> None: + if _is_postgresql(): + with op.get_context().autocommit_block(): + for index_name, columns in _INDEXES: + op.create_index(index_name, _TABLE_NAME, columns, postgresql_concurrently=True) + return + for index_name, columns in _INDEXES: + op.create_index(index_name, _TABLE_NAME, columns) + + +def downgrade() -> None: + if _is_postgresql(): + with op.get_context().autocommit_block(): + for index_name, _ in reversed(_INDEXES): + op.drop_index(index_name, table_name=_TABLE_NAME, postgresql_concurrently=True) + return + for index_name, _ in reversed(_INDEXES): + op.drop_index(index_name, table_name=_TABLE_NAME) diff --git a/api/models/workflow.py b/api/models/workflow.py index d002d5eec7f..58135e79334 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -1476,6 +1476,8 @@ class WorkflowRunArchiveBundle(DefaultFieldsDCMixin, TypeBase): name="workflow_run_archive_bundle_identity_uq", ), sa.Index("workflow_run_archive_bundle_tenant_month_idx", "tenant_id", "year", "month"), + sa.Index("workflow_run_archive_bundle_month_id_idx", "year", "month", "id"), + sa.Index("workflow_run_archive_bundle_month_shard_id_idx", "year", "month", "shard", "id"), ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) diff --git a/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py b/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py index 31435d2f7ea..eb69b198202 100644 --- a/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py +++ b/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py @@ -6,7 +6,10 @@ This service archives workflow run logs for paid plan users older than the confi Archive V2 writes bundle-level Parquet objects. A bundle contains many workflow runs and their related table rows. Bundle metadata lives in the object-store manifest as the recoverable source of truth. Completed bundles are also -mirrored into a small database index so console listing and download jobs do not list object storage online. +published into a small database catalog so console listing, download, and maintenance jobs do not list object storage +online. Archive success requires that catalog publication to commit. A retry reconciles only its known manifest key +against an already-published shard index; a missing shard index fails closed rather than rebuilding it by scanning +historical manifests. Archive campaigns should use fixed absolute UTC windows for every tenant-prefix/shard execution. Relative windows are evaluated at process start and are not safe for multi-day rollout because each command would scan a different window. @@ -639,7 +642,6 @@ class WorkflowRunArchiver: if storage is None: raise ArchiveStorageNotConfiguredError("Archive storage not configured") if storage.object_exists(self._get_manifest_object_key(identity)): - self._write_bundle_index(storage, identity) self._sync_existing_bundle_index(session, storage, identity) result.success = True result.skipped = True @@ -651,6 +653,8 @@ class WorkflowRunArchiver: runs = [run for run in runs if run.id not in archived_run_ids] result.skipped_run_count = original_run_count - len(runs) if not runs: + # Historical catalog rows are a rollout precondition. New bundles commit their catalog row before + # publishing the shard index, so this idempotency path never needs to rescan prior manifests. result.run_count = 0 result.success = True result.skipped = True @@ -663,7 +667,6 @@ class WorkflowRunArchiver: result.object_prefix = identity.object_prefix result.run_count = len(runs) if storage.object_exists(self._get_manifest_object_key(identity)): - self._write_bundle_index(storage, identity) self._sync_existing_bundle_index(session, storage, identity) result.success = True result.skipped = True @@ -695,10 +698,10 @@ class WorkflowRunArchiver: for table_name, payload in table_payloads.items(): storage.put_object(self._get_table_object_key(identity, table_name), payload) storage.put_object(self._get_manifest_object_key(identity), manifest_data) - self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs]) manifest = decode_archive_bundle_manifest(manifest_data) upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data)) session.commit() + self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs]) logger.info( "Archived workflow run bundle %s: tenant=%s runs=%s tables=%s object_prefix=%s", @@ -730,16 +733,18 @@ class WorkflowRunArchiver: storage: ArchiveStorage, identity: ArchiveBundleIdentity, ) -> None: - """Best-effort DB index sync for a bundle whose manifest already exists in archive storage.""" + """Publish a known manifest to the DB catalog and reconcile its existing shard index.""" manifest_key = self._get_manifest_object_key(identity) - try: - manifest_data = storage.get_object(manifest_key) - manifest = decode_archive_bundle_manifest(manifest_data) - upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data)) - session.commit() - except Exception: - session.rollback() - logger.warning("Failed to sync workflow archive bundle index for %s", manifest_key, exc_info=True) + manifest_data = storage.get_object(manifest_key) + manifest = decode_archive_bundle_manifest(manifest_data) + upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data)) + session.commit() + self._merge_bundle_manifest_into_index( + storage, + identity, + manifest["run_ids"], + require_existing_index=True, + ) def _lock_runs_for_archive( self, @@ -1029,10 +1034,20 @@ class WorkflowRunArchiver: storage: ArchiveStorage, identity: ArchiveBundleIdentity, run_ids: Sequence[str], + *, + require_existing_index: bool = False, ) -> ArchiveBundleIndexDict: + """ + Merge one bundle into its shard index. + + Retries for a known manifest set ``require_existing_index`` so they never create a partial index by scanning + or overwriting a shard whose historical entries cannot be proven from this one manifest. + """ index_key = self._get_index_object_key(identity) if storage.object_exists(index_key): index = self._load_bundle_index(storage, identity) + elif require_existing_index: + raise RuntimeError(f"archive shard index missing while reconciling known manifest: {index_key}") else: index = self._build_bundle_index(storage, identity) diff --git a/api/services/retention/workflow_run/bundle_archive_maintenance.py b/api/services/retention/workflow_run/bundle_archive_maintenance.py index 8d90a1d4ac7..54a3d293a76 100644 --- a/api/services/retention/workflow_run/bundle_archive_maintenance.py +++ b/api/services/retention/workflow_run/bundle_archive_maintenance.py @@ -1,14 +1,16 @@ """ Maintain V2 workflow-run archive bundles. -Archive V2 keeps object-store manifests as the recoverable bundle source of truth. This maintenance module still -discovers delete/restore targets by listing `manifest.json` objects and uses object-store marker files for -delete/restore state. The separate database bundle index is intended for console listing and download jobs, not as the -source of truth for destructive maintenance. +Archive V2 keeps object-store manifests as the recoverable bundle source of truth. Delete and restore discover a +bounded page of candidates from `workflow_run_archive_bundles`, optionally restricted to one exact archive shard, then +construct each immutable manifest key from the catalog identity. They never list the object-store namespace. +Object-store marker files keep delete/restore idempotent, while the caller persists a non-dry-run cursor only after a +candidate succeeds. Each bundle is processed in its own database transaction. A failed bundle leaves source rows unchanged unless the transaction has already committed; marker handling makes the next run able to reconcile the common committed-but-marker -not-updated case. +not-updated case. Restore never skips a bundle with a missing deleted marker when deletion started or source rows have +drifted, so an external cursor cannot pass an interrupted delete. """ import datetime @@ -38,6 +40,7 @@ from models.workflow import ( WorkflowPause, WorkflowPauseReason, WorkflowRun, + WorkflowRunArchiveBundle, ) from services.retention.workflow_run.constants import ( ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME, @@ -51,7 +54,6 @@ from services.retention.workflow_run.constants import ( logger = logging.getLogger(__name__) -_ARCHIVE_ROOT_PREFIX = "workflow-runs/v2/" _CHUNK_SIZE = 5_000 @@ -84,11 +86,28 @@ class BundleManifest(TypedDict): @dataclass(frozen=True) -class BundleReference: - """Object-store reference for one V2 archive bundle.""" +class ArchiveBundleCatalogEntry: + """Immutable catalog identity and manifest-derived metrics for one V2 archive bundle.""" + catalog_id: str + tenant_id: str + year: int + month: int + shard: str + bundle_id: str + workflow_run_count: int + row_count: int + archive_bytes: int + + +@dataclass(frozen=True) +class BundleReference: + """Verified object-store reference for one catalog candidate.""" + + catalog: ArchiveBundleCatalogEntry object_prefix: str manifest_key: str + manifest_size_bytes: int manifest: BundleManifest @@ -96,6 +115,7 @@ class BundleReference: class BundleOperationResult: """Result for one V2 bundle delete or restore operation.""" + catalog_id: str bundle_id: str tenant_id: str object_prefix: str @@ -128,6 +148,8 @@ class BundleOperationSummary: archive_bytes: int = 0 elapsed_time: float = 0.0 validation_time: float = 0.0 + next_catalog_id: str | None = None + preview_next_catalog_id: str | None = None table_counts: dict[str, int] = field(default_factory=dict) results: list[BundleOperationResult] = field(default_factory=list) @@ -180,162 +202,334 @@ RESTORE_ORDER = [ "workflow_trigger_logs", ] +DELETE_ORDER = [ + "workflow_pause_reasons", + "workflow_node_execution_offload", + "workflow_trigger_logs", + "workflow_app_logs", + "workflow_node_executions", + "workflow_pauses", + "workflow_runs", +] + class WorkflowRunBundleArchiveMaintenance: """ Delete and restore V2 workflow-run archive bundles. + Delete accepts already-missing source rows only when every remaining row in the bundle scope is an unchanged + archive subset. It then removes only those verified primary keys, so unrelated or unarchived rows fail closed. + Non-dry-run delete and restore serialize on the existing archive catalog row before checking markers or changing + source rows. + Args: dry_run: Validate and report counts without changing source rows or object-store markers. - strict_content_validation: Compare source-table content checksums against Parquet content before destructive - delete and after restore. Keep enabled for real maintenance. - stop_on_error: Stop batch processing after the first failed bundle. + strict_content_validation: Compare restored source-table content checksums against Parquet content. Delete + always validates that every remaining live row belongs to and matches the archive before removing it. + storage: Optional archive storage implementation. Tests may provide an in-memory implementation. + session_factory: Optional session factory. Each candidate is processed in its own transaction. + + Batches stop at the first error so a returned cursor cannot pass an unhandled candidate. """ dry_run: bool strict_content_validation: bool - stop_on_error: bool + storage: ArchiveStorage | None + session_factory: sessionmaker[Session] def __init__( self, *, dry_run: bool = False, strict_content_validation: bool = True, - stop_on_error: bool = True, + storage: ArchiveStorage | None = None, + session_factory: sessionmaker[Session] | None = None, ) -> None: self.dry_run = dry_run self.strict_content_validation = strict_content_validation - self.stop_on_error = stop_on_error + self.storage = storage + self.session_factory = session_factory or sessionmaker(bind=db.engine, expire_on_commit=False) def delete_batch( self, *, tenant_ids: Sequence[str] | None, - start_date: datetime.datetime, - end_date: datetime.datetime, + target_year: int, + target_month: int, + after_catalog_id: str | None, limit: int, + shard: str | None = None, ) -> BundleOperationSummary: - """Validate and delete source rows for archived V2 bundles in the requested created_at window.""" + """Validate and delete one keyset page, optionally scoped to an exact archive shard.""" return self._process_batch( operation="delete", tenant_ids=tenant_ids, - start_date=start_date, - end_date=end_date, + target_year=target_year, + target_month=target_month, + after_catalog_id=after_catalog_id, limit=limit, + shard=shard, ) def restore_batch( self, *, tenant_ids: Sequence[str] | None, - start_date: datetime.datetime, - end_date: datetime.datetime, + target_year: int, + target_month: int, + after_catalog_id: str | None, limit: int, ) -> BundleOperationSummary: - """Restore source rows for deleted V2 bundles in the requested created_at window.""" + """Restore source rows for a keyset page of deleted V2 bundles in one calendar month.""" return self._process_batch( operation="restore", tenant_ids=tenant_ids, - start_date=start_date, - end_date=end_date, + target_year=target_year, + target_month=target_month, + after_catalog_id=after_catalog_id, limit=limit, + shard=None, ) + def validate_catalog_shards( + self, + *, + target_year: int, + target_month: int, + shard_total: int, + tenant_ids: Sequence[str] | None = None, + ) -> None: + """ + Fail before a parallel delete when the requested closed-month scope contains a different shard layout. + + A subset of the expected shards is valid because an archive shard may legitimately contain no bundles. Any + other shard name indicates a historical or mixed-layout month that must be handled by the serial delete path. + """ + if not 1 <= shard_total <= 16: + raise ValueError("shard_total must be between 1 and 16") + expected_shards = tuple(f"{index:02d}-of-{shard_total:02d}" for index in range(shard_total)) + conditions = [ + WorkflowRunArchiveBundle.year == target_year, + WorkflowRunArchiveBundle.month == target_month, + ] + if tenant_ids is not None: + conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids)) + statement = ( + select(WorkflowRunArchiveBundle.shard) + .where( + *conditions, + WorkflowRunArchiveBundle.shard.not_in(expected_shards), + ) + .distinct() + .order_by(WorkflowRunArchiveBundle.shard.asc()) + ) + with self.session_factory() as session: + unexpected_shards = list(session.scalars(statement)) + if unexpected_shards: + raise ValueError( + "archive catalog month contains unexpected shards for " + f"{shard_total}-way delete: {', '.join(unexpected_shards)}" + ) + def _process_batch( self, *, operation: str, tenant_ids: Sequence[str] | None, - start_date: datetime.datetime, - end_date: datetime.datetime, + target_year: int, + target_month: int, + after_catalog_id: str | None, limit: int, + shard: str | None, ) -> BundleOperationSummary: start_time = time.time() summary = BundleOperationSummary(operation=operation) if tenant_ids is not None and not tenant_ids: return summary - storage = self._get_archive_storage() - bundle_refs = self._list_bundle_refs( - storage, - operation=operation, + storage = self.storage or self._get_archive_storage() + catalog_entries = self._list_catalog_entries( tenant_ids=tenant_ids, - start_date=start_date, - end_date=end_date, + target_year=target_year, + target_month=target_month, + after_catalog_id=after_catalog_id, limit=limit, + shard=shard, ) - logger.info("Found %s V2 archive bundles for %s", len(bundle_refs), operation) - session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - for bundle_ref in bundle_refs: - with session_maker() as session: - if operation == "delete": - result = self._delete_bundle(session, storage, bundle_ref) - elif operation == "restore": - result = self._restore_bundle(session, storage, bundle_ref) - else: - raise ValueError(f"Unsupported operation: {operation}") + logger.info( + "Found %s V2 archive catalog candidates for %s: year=%s month=%s shard=%s after_catalog_id=%s", + len(catalog_entries), + operation, + target_year, + target_month, + shard, + after_catalog_id, + ) + for catalog_entry in catalog_entries: + try: + bundle_ref = self._build_bundle_reference(storage, catalog_entry) + with self.session_factory() as session: + if operation == "delete": + result = self._delete_bundle(session, storage, bundle_ref) + elif operation == "restore": + result = self._restore_bundle(session, storage, bundle_ref) + else: + raise ValueError(f"Unsupported operation: {operation}") + except Exception as exc: + result = self._new_result_from_catalog_entry(catalog_entry) + result.error = str(exc) + logger.exception( + "Failed to prepare V2 archive bundle %s from catalog %s", + catalog_entry.bundle_id, + catalog_entry.catalog_id, + ) self._merge_result(summary, result) - if not result.success and self.stop_on_error: + if result.success: + if self.dry_run: + summary.preview_next_catalog_id = catalog_entry.catalog_id + else: + summary.next_catalog_id = catalog_entry.catalog_id + else: logger.error("Stopping V2 bundle %s after failure: %s", operation, result.error) break summary.elapsed_time = time.time() - start_time return summary - def _list_bundle_refs( + def _list_catalog_entries( self, - storage: ArchiveStorage, *, - operation: str, tenant_ids: Sequence[str] | None, - start_date: datetime.datetime, - end_date: datetime.datetime, + target_year: int, + target_month: int, + after_catalog_id: str | None, limit: int, - ) -> list[BundleReference]: - start_date = self._to_naive_utc(start_date) - end_date = self._to_naive_utc(end_date) - manifest_keys = self._list_manifest_keys(storage, tenant_ids) - refs: list[BundleReference] = [] - for manifest_key in manifest_keys: - manifest_data = self._get_checked_object(storage, manifest_key) - object_prefix = manifest_key.removesuffix(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}") - manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix) - min_created_at = self._parse_manifest_datetime(manifest["min_created_at"]) - max_created_at = self._parse_manifest_datetime(manifest["max_created_at"]) - if max_created_at < start_date or min_created_at >= end_date: - continue - if tenant_ids and manifest["tenant_id"] not in tenant_ids: - continue - if operation == "delete" and self._is_deleted(storage, object_prefix): - continue - if operation == "restore" and not self._is_deleted(storage, object_prefix): - continue - refs.append(BundleReference(object_prefix=object_prefix, manifest_key=manifest_key, manifest=manifest)) + shard: str | None = None, + ) -> list[ArchiveBundleCatalogEntry]: + """Read one bounded, stable-order candidate page from the database catalog.""" + conditions = [ + WorkflowRunArchiveBundle.year == target_year, + WorkflowRunArchiveBundle.month == target_month, + ] + if tenant_ids: + conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids)) + if shard is not None: + conditions.append(WorkflowRunArchiveBundle.shard == shard) + if after_catalog_id: + conditions.append(WorkflowRunArchiveBundle.id > after_catalog_id) - refs.sort( - key=lambda ref: ( - self._parse_manifest_datetime(ref.manifest["min_created_at"]), - ref.manifest["tenant_id"], - ref.manifest["bundle_id"], - ) + statement = ( + select(WorkflowRunArchiveBundle).where(*conditions).order_by(WorkflowRunArchiveBundle.id.asc()).limit(limit) ) - return refs[:limit] + with self.session_factory() as session: + if after_catalog_id: + cursor_bundle = session.get(WorkflowRunArchiveBundle, after_catalog_id) + self._validate_catalog_cursor_scope( + cursor_bundle, + tenant_ids=tenant_ids, + target_year=target_year, + target_month=target_month, + shard=shard, + ) + bundles = list(session.scalars(statement)) + return [ + ArchiveBundleCatalogEntry( + catalog_id=bundle.id, + tenant_id=bundle.tenant_id, + year=bundle.year, + month=bundle.month, + shard=bundle.shard, + bundle_id=bundle.bundle_id, + workflow_run_count=bundle.workflow_run_count, + row_count=bundle.row_count, + archive_bytes=bundle.archive_bytes, + ) + for bundle in bundles + ] @staticmethod - def _list_manifest_keys(storage: ArchiveStorage, tenant_ids: Sequence[str] | None) -> list[str]: - keys: list[str] = [] - if tenant_ids: - prefixes = [ - f"{_ARCHIVE_ROOT_PREFIX}tenant_prefix={tenant_id[0].lower()}/tenant_id={tenant_id}/" - for tenant_id in tenant_ids - ] - else: - prefixes = [_ARCHIVE_ROOT_PREFIX] - for prefix in prefixes: - keys.extend(storage.list_objects(prefix)) - return sorted(key for key in keys if key.endswith(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}")) + def _validate_catalog_cursor_scope( + cursor_bundle: WorkflowRunArchiveBundle | None, + *, + tenant_ids: Sequence[str] | None, + target_year: int, + target_month: int, + shard: str | None, + ) -> None: + """Reject a keyset cursor that cannot safely represent the requested catalog scope.""" + if cursor_bundle is None: + raise ValueError("after_catalog_id does not exist in the workflow run archive bundle catalog") + if cursor_bundle.year != target_year or cursor_bundle.month != target_month: + raise ValueError("after_catalog_id is outside the requested archive month") + if tenant_ids is not None and cursor_bundle.tenant_id not in tenant_ids: + raise ValueError("after_catalog_id is outside the requested tenant scope") + if shard is not None and cursor_bundle.shard != shard: + raise ValueError("after_catalog_id is outside the requested archive shard") + + def _build_bundle_reference( + self, + storage: ArchiveStorage, + catalog_entry: ArchiveBundleCatalogEntry, + ) -> BundleReference: + """Load and bind one manifest to the catalog row that selected it.""" + object_prefix = self._catalog_object_prefix(catalog_entry) + manifest_key = f"{object_prefix}/{ARCHIVE_BUNDLE_MANIFEST_NAME}" + manifest_data = storage.get_object(manifest_key) + manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix) + self._validate_manifest_catalog_identity(manifest, catalog_entry) + return BundleReference( + catalog=catalog_entry, + object_prefix=object_prefix, + manifest_key=manifest_key, + manifest_size_bytes=len(manifest_data), + manifest=manifest, + ) + + @staticmethod + def _catalog_object_prefix(catalog_entry: ArchiveBundleCatalogEntry) -> str: + """Construct the immutable V2 bundle prefix from its database catalog identity.""" + if not catalog_entry.tenant_id: + raise ValueError("archive catalog tenant_id must not be empty") + if not 1 <= catalog_entry.month <= 12: + raise ValueError(f"archive catalog month is invalid: {catalog_entry.month}") + return ( + f"workflow-runs/v2/tenant_prefix={catalog_entry.tenant_id[0].lower()}/" + f"tenant_id={catalog_entry.tenant_id}/year={catalog_entry.year:04d}/" + f"month={catalog_entry.month:02d}/shard={catalog_entry.shard}/bundle={catalog_entry.bundle_id}" + ) + + @staticmethod + def _validate_manifest_catalog_identity( + manifest: BundleManifest, + catalog_entry: ArchiveBundleCatalogEntry, + ) -> None: + """Fail closed when the catalog locator and manifest identify different immutable bundles.""" + expected_identity = ( + catalog_entry.tenant_id, + catalog_entry.year, + catalog_entry.month, + catalog_entry.shard, + catalog_entry.bundle_id, + ) + manifest_identity = ( + manifest["tenant_id"], + manifest["year"], + manifest["month"], + manifest["shard"], + manifest["bundle_id"], + ) + if manifest_identity != expected_identity: + raise ValueError( + f"archive manifest identity does not match catalog: expected={expected_identity}, " + f"actual={manifest_identity}" + ) + if manifest["workflow_run_count"] != catalog_entry.workflow_run_count: + raise ValueError("archive manifest workflow_run_count does not match catalog") + manifest_row_count = sum(table["row_count"] for table in manifest["tables"].values()) + if manifest_row_count != catalog_entry.row_count: + raise ValueError("archive manifest row_count does not match catalog") def _delete_bundle( self, @@ -344,37 +538,61 @@ class WorkflowRunBundleArchiveMaintenance: bundle_ref: BundleReference, ) -> BundleOperationResult: start_time = time.time() - result = self._new_result(bundle_ref.manifest) + result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id) try: validation_start = time.time() + if not self.dry_run: + self._lock_catalog_entry(session, bundle_ref.catalog) + if self._is_restore_started(storage, bundle_ref.object_prefix): + raise ValueError("restore started marker exists; reconcile restore before delete") + + deleted_marker_exists = self._is_deleted(storage, bundle_ref.object_prefix) manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref) result.table_counts = self._manifest_table_counts(manifest) result.archive_bytes = archive_bytes - self._lock_workflow_runs(session, manifest["run_ids"]) - if self._is_delete_started(storage, bundle_ref.object_prefix) and self._live_counts_match( - session, manifest, expected_present=False - ): + live_records = self._load_live_bundle_records( + session, + manifest, + table_records, + lock=not self.dry_run, + ) + if deleted_marker_exists: + live_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES} + if any(live_counts.values()): + raise ValueError(f"Live rows exist for bundle with deleted marker: {live_counts}") + if not self.dry_run: + self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME) + self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME) result.validation_time = time.time() - validation_start + result.success = True + result.elapsed_time = time.time() - start_time + return result + + self._validate_live_archive_subset(manifest, table_records, live_records) + result.validation_time = time.time() - validation_start + + if not any(live_records[table_name] for table_name in ARCHIVED_TABLES): if not self.dry_run: self._mark_deleted(storage, bundle_ref.object_prefix) self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME) + self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME) result.success = True + result.elapsed_time = time.time() - start_time return result - self._validate_live_counts(session, manifest, expected_present=True) - if self.strict_content_validation: - self._validate_live_content(session, table_records) - result.validation_time = time.time() - validation_start - if not self.dry_run: self._put_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME) - deleted_counts = self._delete_bundle_rows(session, table_records) - if deleted_counts != result.table_counts: + expected_deleted_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES} + deleted_counts = self._delete_bundle_rows(session, live_records) + if deleted_counts != expected_deleted_counts: raise ValueError( - f"Deleted row count mismatch: expected={result.table_counts}, actual={deleted_counts}" + f"Deleted row count mismatch: expected={expected_deleted_counts}, actual={deleted_counts}" ) - self._validate_live_counts(session, manifest, expected_present=False) + remaining_records = self._load_live_bundle_records(session, manifest, table_records, lock=True) + remaining_counts = {table_name: len(remaining_records[table_name]) for table_name in ARCHIVED_TABLES} + if any(remaining_counts.values()): + raise ValueError(f"Live rows remain after bundle delete: {remaining_counts}") session.commit() self._mark_deleted(storage, bundle_ref.object_prefix) self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME) @@ -394,9 +612,25 @@ class WorkflowRunBundleArchiveMaintenance: bundle_ref: BundleReference, ) -> BundleOperationResult: start_time = time.time() - result = self._new_result(bundle_ref.manifest) + result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id) try: validation_start = time.time() + if not self.dry_run: + self._lock_catalog_entry(session, bundle_ref.catalog) + if not self._is_deleted(storage, bundle_ref.object_prefix): + # A committed delete may be interrupted before `.deleted` is written. Do not let restore advance its + # cursor over that state: retry delete to reconcile it, or investigate source-row drift first. + if self._is_delete_started(storage, bundle_ref.object_prefix): + raise ValueError("delete started marker exists without a deleted marker; reconcile delete first") + restore_started = self._is_restore_started(storage, bundle_ref.object_prefix) + self._validate_live_counts(session, bundle_ref.manifest, expected_present=True) + result.validation_time = time.time() - validation_start + if restore_started and not self.dry_run: + self._mark_restored(storage, bundle_ref.object_prefix) + result.success = True + result.elapsed_time = time.time() - start_time + return result + manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref) result.table_counts = self._manifest_table_counts(manifest) result.archive_bytes = archive_bytes @@ -408,6 +642,7 @@ class WorkflowRunBundleArchiveMaintenance: if not self.dry_run: self._mark_restored(storage, bundle_ref.object_prefix) result.success = True + result.elapsed_time = time.time() - start_time return result self._validate_live_counts(session, manifest, expected_present=False) @@ -432,13 +667,40 @@ class WorkflowRunBundleArchiveMaintenance: return result @staticmethod - def _new_result(manifest: BundleManifest) -> BundleOperationResult: + def _new_result(manifest: BundleManifest, catalog_id: str) -> BundleOperationResult: return BundleOperationResult( + catalog_id=catalog_id, bundle_id=manifest["bundle_id"], tenant_id=manifest["tenant_id"], object_prefix=manifest["object_prefix"], ) + @staticmethod + def _new_result_from_catalog_entry(catalog_entry: ArchiveBundleCatalogEntry) -> BundleOperationResult: + try: + object_prefix = WorkflowRunBundleArchiveMaintenance._catalog_object_prefix(catalog_entry) + except ValueError: + object_prefix = "" + return BundleOperationResult( + catalog_id=catalog_entry.catalog_id, + bundle_id=catalog_entry.bundle_id, + tenant_id=catalog_entry.tenant_id, + object_prefix=object_prefix, + ) + + @staticmethod + def _lock_catalog_entry(session: Session, catalog_entry: ArchiveBundleCatalogEntry) -> None: + locked_catalog_id = session.scalar( + select(WorkflowRunArchiveBundle.id) + .where( + WorkflowRunArchiveBundle.id == catalog_entry.catalog_id, + WorkflowRunArchiveBundle.tenant_id == catalog_entry.tenant_id, + ) + .with_for_update() + ) + if locked_catalog_id is None: + raise ValueError("archive catalog row disappeared before bundle maintenance") + def _validate_archive_object( self, storage: ArchiveStorage, @@ -446,7 +708,7 @@ class WorkflowRunBundleArchiveMaintenance: ) -> tuple[BundleManifest, dict[str, list[dict[str, Any]]], int]: manifest = bundle_ref.manifest table_records: dict[str, list[dict[str, Any]]] = {} - total_size = len(storage.get_object(bundle_ref.manifest_key)) + total_size = bundle_ref.manifest_size_bytes for table_name in ARCHIVED_TABLES: info = manifest["tables"][table_name] payload = self._get_checked_object(storage, info["object_key"]) @@ -469,12 +731,14 @@ class WorkflowRunBundleArchiveMaintenance: f"expected={info['row_count']}, actual={len(records)}" ) table_records[table_name] = records + if total_size != bundle_ref.catalog.archive_bytes: + raise ValueError( + f"Archive object total size mismatch: expected={bundle_ref.catalog.archive_bytes}, actual={total_size}" + ) return manifest, table_records, total_size @staticmethod def _get_checked_object(storage: ArchiveStorage, object_key: str) -> bytes: - if not storage.object_exists(object_key): - raise FileNotFoundError(f"Archive object not found: {object_key}") return storage.get_object(object_key) @staticmethod @@ -539,6 +803,110 @@ class WorkflowRunBundleArchiveMaintenance: table = pq.read_table(io.BytesIO(payload)) return table.to_pylist() + def _load_live_bundle_records( + self, + session: Session, + manifest: BundleManifest, + table_records: dict[str, list[dict[str, Any]]], + *, + lock: bool, + ) -> dict[str, list[dict[str, Any]]]: + """Load the complete live scope for a bundle, including archived rows whose relationship fields drifted.""" + run_ids = manifest["run_ids"] + archive_ids = { + table_name: [str(record["id"]) for record in table_records[table_name]] for table_name in ARCHIVED_TABLES + } + live_node_ids = self._select_ids_by_run_ids(session, WorkflowNodeExecutionModel, run_ids) + live_pause_ids = self._select_ids_by_run_ids(session, WorkflowPause, run_ids) + node_ids = sorted(set(archive_ids["workflow_node_executions"]) | set(live_node_ids)) + pause_ids = sorted(set(archive_ids["workflow_pauses"]) | set(live_pause_ids)) + + def load_scope( + table_name: str, + model: Any, + scope_column: Any, + scope_ids: Sequence[str], + ) -> list[dict[str, Any]]: + return self._merge_records_by_id( + self._load_records_by_column(session, model, scope_column, scope_ids, lock=lock), + self._load_records_by_column(session, model, model.id, archive_ids[table_name], lock=lock), + ) + + return { + "workflow_pause_reasons": load_scope( + "workflow_pause_reasons", WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids + ), + "workflow_node_execution_offload": load_scope( + "workflow_node_execution_offload", + WorkflowNodeExecutionOffload, + WorkflowNodeExecutionOffload.node_execution_id, + node_ids, + ), + "workflow_trigger_logs": load_scope( + "workflow_trigger_logs", WorkflowTriggerLog, WorkflowTriggerLog.workflow_run_id, run_ids + ), + "workflow_app_logs": load_scope( + "workflow_app_logs", WorkflowAppLog, WorkflowAppLog.workflow_run_id, run_ids + ), + "workflow_node_executions": load_scope( + "workflow_node_executions", + WorkflowNodeExecutionModel, + WorkflowNodeExecutionModel.workflow_run_id, + run_ids, + ), + "workflow_pauses": load_scope("workflow_pauses", WorkflowPause, WorkflowPause.workflow_run_id, run_ids), + "workflow_runs": self._load_records_by_column( + session, + WorkflowRun, + WorkflowRun.id, + sorted(set(run_ids) | set(archive_ids["workflow_runs"])), + lock=lock, + ), + } + + @classmethod + def _validate_live_archive_subset( + cls, + manifest: BundleManifest, + table_records: dict[str, list[dict[str, Any]]], + live_records: dict[str, list[dict[str, Any]]], + ) -> None: + """Require every live row in the bundle scope to exist unchanged in the validated archive.""" + manifest_run_ids = {str(run_id) for run_id in manifest["run_ids"]} + if len(manifest_run_ids) != len(manifest["run_ids"]): + raise ValueError("archive manifest contains duplicate workflow run IDs") + + archive_records_by_id: dict[str, dict[str, dict[str, Any]]] = {} + for table_name in ARCHIVED_TABLES: + records_by_id = {str(record["id"]): record for record in table_records[table_name]} + if len(records_by_id) != len(table_records[table_name]): + raise ValueError(f"archive contains duplicate row IDs for {table_name}") + archive_records_by_id[table_name] = records_by_id + + if set(archive_records_by_id["workflow_runs"]) != manifest_run_ids: + raise ValueError("archive workflow run IDs do not match manifest run_ids") + + for table_name in ARCHIVED_TABLES: + live_ids = [str(record["id"]) for record in live_records[table_name]] + if len(set(live_ids)) != len(live_ids): + raise ValueError(f"live scope contains duplicate row IDs for {table_name}") + + archive_by_id = archive_records_by_id[table_name] + extra_ids = sorted(set(live_ids) - set(archive_by_id)) + if extra_ids: + raise ValueError( + f"Live bundle scope contains rows missing from archive for {table_name}: {extra_ids[:10]}" + ) + + archive_subset = [archive_by_id[row_id] for row_id in live_ids] + live_checksum = cls._records_checksum(live_records[table_name]) + archive_checksum = cls._records_checksum(archive_subset) + if live_checksum != archive_checksum: + raise ValueError( + f"Live/archive subset content checksum mismatch for {table_name}: " + f"expected={archive_checksum}, actual={live_checksum}" + ) + def _validate_live_counts( self, session: Session, @@ -617,26 +985,13 @@ class WorkflowRunBundleArchiveMaintenance: def _delete_bundle_rows( self, session: Session, - table_records: dict[str, list[dict[str, Any]]], + live_records: dict[str, list[dict[str, Any]]], ) -> dict[str, int]: - run_ids = [str(record["id"]) for record in table_records["workflow_runs"]] - node_ids = [str(record["id"]) for record in table_records["workflow_node_executions"]] - pause_ids = [str(record["id"]) for record in table_records["workflow_pauses"]] - deleted_counts = dict.fromkeys(ARCHIVED_TABLES, 0) - deleted_counts["workflow_pause_reasons"] = self._delete_by_column( - session, WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids - ) - deleted_counts["workflow_node_execution_offload"] = self._delete_by_column( - session, WorkflowNodeExecutionOffload, WorkflowNodeExecutionOffload.node_execution_id, node_ids - ) - deleted_counts["workflow_trigger_logs"] = self._delete_by_run_ids(session, WorkflowTriggerLog, run_ids) - deleted_counts["workflow_app_logs"] = self._delete_by_run_ids(session, WorkflowAppLog, run_ids) - deleted_counts["workflow_node_executions"] = self._delete_by_run_ids( - session, WorkflowNodeExecutionModel, run_ids - ) - deleted_counts["workflow_pauses"] = self._delete_by_run_ids(session, WorkflowPause, run_ids) - deleted_counts["workflow_runs"] = self._delete_by_run_ids(session, WorkflowRun, run_ids) + for table_name in DELETE_ORDER: + model = TABLE_MODELS[table_name] + row_ids = [str(record["id"]) for record in live_records[table_name]] + deleted_counts[table_name] = self._delete_by_column(session, model, model.id, row_ids) return deleted_counts def _restore_bundle_rows( @@ -708,11 +1063,6 @@ class WorkflowRunBundleArchiveMaintenance: payload = json.dumps(normalized, sort_keys=True, default=str, ensure_ascii=False, separators=(",", ":")) return ArchiveStorage.compute_checksum(payload.encode("utf-8")) - @staticmethod - def _lock_workflow_runs(session: Session, run_ids: Sequence[str]) -> None: - for chunk in WorkflowRunBundleArchiveMaintenance._chunks(run_ids, _CHUNK_SIZE): - list(session.scalars(select(WorkflowRun.id).where(WorkflowRun.id.in_(chunk)).with_for_update())) - @staticmethod def _select_ids_by_run_ids( session: Session, @@ -757,8 +1107,16 @@ class WorkflowRunBundleArchiveMaintenance: session: Session, model: Any, run_ids: Sequence[str], + *, + lock: bool = False, ) -> list[dict[str, Any]]: - return self._load_records_by_column(session, model, self._run_id_column(model), run_ids) + return self._load_records_by_column( + session, + model, + self._run_id_column(model), + run_ids, + lock=lock, + ) def _load_records_by_column( self, @@ -766,23 +1124,23 @@ class WorkflowRunBundleArchiveMaintenance: model: Any, column: Any, values: Sequence[str], + *, + lock: bool = False, ) -> list[dict[str, Any]]: if not values: return [] rows: list[Any] = [] for chunk in self._chunks(values, _CHUNK_SIZE): - rows.extend(session.scalars(select(model).where(column.in_(chunk)))) + statement = select(model).where(column.in_(chunk)).order_by(model.id.asc()) + if lock: + statement = statement.with_for_update() + rows.extend(session.scalars(statement)) return [self._row_to_dict(row) for row in rows] @staticmethod - def _delete_by_run_ids( - session: Session, - model: Any, - run_ids: Sequence[str], - ) -> int: - return WorkflowRunBundleArchiveMaintenance._delete_by_column( - session, model, WorkflowRunBundleArchiveMaintenance._run_id_column(model), run_ids - ) + def _merge_records_by_id(*record_groups: list[dict[str, Any]]) -> list[dict[str, Any]]: + records_by_id = {str(record["id"]): record for record_group in record_groups for record in record_group} + return [records_by_id[row_id] for row_id in sorted(records_by_id)] @staticmethod def _run_id_column(model: Any) -> Any: @@ -813,6 +1171,10 @@ class WorkflowRunBundleArchiveMaintenance: def _is_delete_started(storage: ArchiveStorage, object_prefix: str) -> bool: return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME}") + @staticmethod + def _is_restore_started(storage: ArchiveStorage, object_prefix: str) -> bool: + return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME}") + @staticmethod def _mark_deleted(storage: ArchiveStorage, object_prefix: str) -> None: WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME) @@ -820,10 +1182,13 @@ class WorkflowRunBundleArchiveMaintenance: @staticmethod def _mark_restored(storage: ArchiveStorage, object_prefix: str) -> None: WorkflowRunBundleArchiveMaintenance._delete_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME) + WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME) + WorkflowRunBundleArchiveMaintenance._delete_marker( + storage, object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME + ) WorkflowRunBundleArchiveMaintenance._delete_marker( storage, object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME ) - WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME) @staticmethod def _put_marker(storage: ArchiveStorage, object_prefix: str, marker_name: str) -> None: @@ -836,16 +1201,6 @@ class WorkflowRunBundleArchiveMaintenance: if storage.object_exists(marker_key): storage.delete_object(marker_key) - @staticmethod - def _parse_manifest_datetime(value: str) -> datetime.datetime: - return WorkflowRunBundleArchiveMaintenance._to_naive_utc(datetime.datetime.fromisoformat(value)) - - @staticmethod - def _to_naive_utc(value: datetime.datetime) -> datetime.datetime: - if value.tzinfo is None: - return value - return value.astimezone(datetime.UTC).replace(tzinfo=None) - @staticmethod def _chunks(values: Sequence[Any], size: int) -> list[Sequence[Any]]: return [values[index : index + size] for index in range(0, len(values), size)] diff --git a/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py b/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py index 0269e99779b..90639ac60d6 100644 --- a/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py +++ b/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py @@ -513,11 +513,94 @@ class TestArchiveRunIdempotency: storage = MagicMock() storage.object_exists.return_value = True - result = archiver._archive_bundle(MagicMock(), storage, [run]) + with patch.object(archiver, "_sync_existing_bundle_index") as sync_existing_bundle_index: + result = archiver._archive_bundle(MagicMock(), storage, [run]) assert result.success is True assert result.skipped is True assert result.error == "bundle already archived" + sync_existing_bundle_index.assert_called_once() + + def test_existing_bundle_catalog_publication_failure_is_not_success(self): + archiver = WorkflowRunArchiver(days=90) + run = _run() + session = MagicMock() + storage = MagicMock() + storage.object_exists.return_value = True + + with patch.object(archiver, "_sync_existing_bundle_index", side_effect=RuntimeError("catalog unavailable")): + result = archiver._archive_bundle(session, storage, [run]) + + assert result.success is False + assert result.error == "catalog unavailable" + session.rollback.assert_called_once() + + def test_retry_repairs_index_after_catalog_commit_then_index_write_failure(self): + archiver = WorkflowRunArchiver(days=90) + run = _run() + identity = archiver._build_bundle_identity([run]) + index_key = archiver._get_index_object_key(identity) + manifest_key = archiver._get_manifest_object_key(identity) + storage = FakeArchiveStorage() + original_put_object = storage.put_object + index_write_count = 0 + + def put_object(key: str, data: bytes) -> str: + nonlocal index_write_count + if key == index_key: + index_write_count += 1 + if index_write_count == 2: + raise RuntimeError("index write failed") + return original_put_object(key, data) + + storage.put_object = MagicMock(side_effect=put_object) + first_session = MagicMock() + first_session.scalar.return_value = None + table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), + patch.object(archiver, "_extract_bundle_data", return_value=table_data), + ): + first_result = archiver._archive_bundle(first_session, storage, [run]) + + assert first_result.success is False + assert first_result.error == "index write failed" + assert manifest_key in storage.objects + assert json.loads(storage.objects[index_key])["run_ids"] == [] + + storage.list_objects = MagicMock(wraps=storage.list_objects) + retry_session = MagicMock() + retry_session.scalar.return_value = None + + retry_result = archiver._archive_bundle(retry_session, storage, [run]) + + assert retry_result.success is True + assert retry_result.skipped is True + assert json.loads(storage.objects[index_key])["manifest_keys"] == [manifest_key] + assert json.loads(storage.objects[index_key])["run_ids"] == [run.id] + storage.list_objects.assert_not_called() + + def test_existing_manifest_with_missing_index_fails_without_partial_rebuild(self): + archiver = WorkflowRunArchiver(days=90) + run = _run() + identity = archiver._build_bundle_identity([run]) + _, _, manifest_data = archiver._build_archive_payload( + identity, + [run], + {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}, + ) + manifest_key = archiver._get_manifest_object_key(identity) + index_key = archiver._get_index_object_key(identity) + storage = FakeArchiveStorage({manifest_key: manifest_data}) + storage.list_objects = MagicMock(wraps=storage.list_objects) + + result = archiver._archive_bundle(MagicMock(), storage, [run]) + + assert result.success is False + assert "archive shard index missing" in (result.error or "") + assert index_key not in storage.objects + storage.list_objects.assert_not_called() def test_successful_bundle_persists_archive_index(self): archiver = WorkflowRunArchiver(days=90) @@ -550,6 +633,28 @@ class TestArchiveRunIdempotency: assert archived_bundle.row_count == 2 session.commit.assert_called_once() + def test_new_bundle_catalog_commit_failure_is_not_success(self): + archiver = WorkflowRunArchiver(days=90) + run = _run(str(uuid.uuid4())) + run.tenant_id = str(uuid.uuid4()) + session = MagicMock() + session.scalar.return_value = None + session.commit.side_effect = RuntimeError("catalog commit failed") + storage = MagicMock() + storage.object_exists.return_value = False + storage.list_objects.return_value = [] + table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), + patch.object(archiver, "_extract_bundle_data", return_value=table_data), + ): + result = archiver._archive_bundle(session, storage, [run]) + + assert result.success is False + assert result.error == "catalog commit failed" + session.rollback.assert_called_once() + def test_index_skips_all_already_archived_runs(self): archiver = WorkflowRunArchiver(days=90) run = MagicMock() diff --git a/api/tests/unit_tests/commands/test_archive_workflow_runs.py b/api/tests/unit_tests/commands/test_archive_workflow_runs.py index 2e2b5df8029..60766b71ca8 100644 --- a/api/tests/unit_tests/commands/test_archive_workflow_runs.py +++ b/api/tests/unit_tests/commands/test_archive_workflow_runs.py @@ -3,9 +3,19 @@ from unittest.mock import MagicMock import click import pytest +from click.testing import CliRunner from sqlalchemy.exc import OperationalError from commands import retention +from services.retention.workflow_run import bundle_archive_maintenance +from services.retention.workflow_run.bundle_archive_maintenance import ( + BundleOperationResult, + BundleOperationSummary, +) + +_CURSOR_0 = "00000000-0000-0000-0000-000000000000" +_CURSOR_1 = "00000000-0000-0000-0000-000000000001" +_CURSOR_2 = "00000000-0000-0000-0000-000000000002" def _db_disconnect_error() -> OperationalError: @@ -24,6 +34,58 @@ def _session_context(session): return context +def _delete_summary( + *, + processed: int, + succeeded: int = 0, + failed: int = 0, + next_catalog_id: str | None = None, + preview_next_catalog_id: str | None = None, + results: list[BundleOperationResult] | None = None, +) -> BundleOperationSummary: + return BundleOperationSummary( + operation="delete", + bundles_processed=processed, + bundles_succeeded=succeeded, + bundles_failed=failed, + next_catalog_id=next_catalog_id, + preview_next_catalog_id=preview_next_catalog_id, + results=results or [], + ) + + +def _patch_bundle_deleter( + monkeypatch: pytest.MonkeyPatch, + summaries: list[BundleOperationSummary], +) -> MagicMock: + deleter = MagicMock() + deleter.delete_batch.side_effect = summaries + monkeypatch.setattr( + bundle_archive_maintenance, + "WorkflowRunBundleArchiveMaintenance", + MagicMock(return_value=deleter), + ) + return deleter + + +@pytest.mark.parametrize( + "command", + [retention.restore_workflow_runs, retention.delete_archived_workflow_runs], +) +def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command): + result = CliRunner().invoke( + command, + ["--tenant-ids", "", "--target-month", "2025-03"], + ) + + assert result.exit_code == 2 + assert "tenant-ids must not be empty" in result.output + + +def test_archive_tenant_id_parser_keeps_omitted_scope_unset(): + assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None + + def test_resolve_archive_tenant_ids_from_plan_uses_explicit_sessions(monkeypatch): end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC) sessions = [MagicMock(name="session-a"), MagicMock(name="session-b")] @@ -137,3 +199,265 @@ def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(mon dry_run=True, delete_after_archive=False, ) + + +def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages(monkeypatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--limit", "2"], + ) + + assert result.exit_code == 0 + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["target_year"] == 2025 + assert deleter.delete_batch.call_args.kwargs["target_month"] == 3 + assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] is None + assert deleter.delete_batch.call_args.kwargs["limit"] == 2 + + +def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), + _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_2), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--limit", "2"], + ) + + assert result.exit_code == 0 + assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ + None, + _CURSOR_1, + _CURSOR_2, + ] + + +def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page(monkeypatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--limit", "2"], + ) + + assert result.exit_code == 0 + assert deleter.delete_batch.call_count == 2 + assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1 + + +def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch): + failed_result = BundleOperationResult( + catalog_id=_CURSOR_2, + bundle_id="bundle-failed", + tenant_id="tenant-1", + object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", + error="archive checksum mismatch", + ) + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_1), + _delete_summary(processed=1, failed=1, results=[failed_result]), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages"], + ) + + assert result.exit_code == 1 + assert deleter.delete_batch.call_count == 2 + assert "target_month=2025-03" in result.output + assert f"failed_catalog_id={_CURSOR_2}" in result.output + assert f"resume_after_catalog_id={_CURSOR_1}" in result.output + + +def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance(monkeypatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [_delete_summary(processed=1, succeeded=1, next_catalog_id=None)], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages"], + ) + + assert result.exit_code == 1 + deleter.delete_batch.assert_called_once() + assert "cursor did not advance" in result.output.lower() + + +def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=1, succeeded=1, preview_next_catalog_id=_CURSOR_1), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--dry-run"], + ) + + assert result.exit_code == 0 + assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ + None, + _CURSOR_1, + ] + + +def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors(monkeypatch): + failed_result = BundleOperationResult( + catalog_id=_CURSOR_2, + bundle_id="bundle-failed", + tenant_id="tenant-1", + object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", + error="archive checksum mismatch", + ) + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary( + processed=1, + succeeded=1, + preview_next_catalog_id=_CURSOR_1, + ), + _delete_summary( + processed=1, + failed=1, + results=[failed_result], + ), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--after-catalog-id", + _CURSOR_0, + "--all-pages", + "--dry-run", + ], + ) + + assert result.exit_code == 1 + assert deleter.delete_batch.call_count == 2 + assert f"failed_catalog_id={_CURSOR_2}" in result.output + assert f"preview_after_catalog_id={_CURSOR_1}" in result.output + assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output + + +def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--after-catalog-id", + _CURSOR_0, + "--all-pages", + ], + ) + + assert result.exit_code == 0 + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] == _CURSOR_0 + + +@pytest.mark.parametrize( + "shard_args", + [ + ["--run-shard-index", "0"], + ["--run-shard-total", "16"], + ["--run-shard-index", "16", "--run-shard-total", "16"], + ["--run-shard-index", "-1", "--run-shard-total", "16"], + ["--run-shard-index", "0", "--run-shard-total", "0"], + ["--run-shard-index", "0", "--run-shard-total", "17"], + ], +) +def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", *shard_args], + ) + + assert result.exit_code == 2 + deleter.delete_batch.assert_not_called() + + +def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--tenant-ids", + "tenant-1", + "--run-shard-index", + "3", + "--run-shard-total", + "16", + ], + ) + + assert result.exit_code == 0 + deleter.validate_catalog_shards.assert_called_once_with( + target_year=2025, + target_month=3, + shard_total=16, + tenant_ids=["tenant-1"], + ) + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16" + + +def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01") + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--run-shard-index", + "3", + "--run-shard-total", + "16", + ], + ) + + assert result.exit_code == 1 + assert "shard preflight failed" in result.output.lower() + assert "00-of-01" in result.output + deleter.delete_batch.assert_not_called() diff --git a/api/tests/unit_tests/libs/test_archive_storage.py b/api/tests/unit_tests/libs/test_archive_storage.py index 659f3cdc6a6..f42bc63d5f4 100644 --- a/api/tests/unit_tests/libs/test_archive_storage.py +++ b/api/tests/unit_tests/libs/test_archive_storage.py @@ -4,7 +4,7 @@ from datetime import datetime from unittest.mock import ANY, MagicMock import pytest -from botocore.exceptions import ClientError +from botocore.exceptions import ClientError, EndpointConnectionError from libs import archive_storage as storage_module from libs.archive_storage import ( @@ -34,6 +34,10 @@ def _client_error(code: str) -> ClientError: return ClientError({"Error": {"Code": code}}, "Operation") +def _network_error() -> EndpointConnectionError: + return EndpointConnectionError(endpoint_url="https://storage.example.com") + + def _mock_client(monkeypatch: pytest.MonkeyPatch): client = MagicMock() client.head_bucket.return_value = None @@ -153,16 +157,38 @@ def test_get_object_returns_bytes(monkeypatch: pytest.MonkeyPatch): assert storage.get_object("key") == b"payload" -def test_get_object_missing(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"]) +def test_get_object_missing(monkeypatch: pytest.MonkeyPatch, error_code: str): _configure_storage(monkeypatch) client, _ = _mock_client(monkeypatch) - client.get_object.side_effect = _client_error("NoSuchKey") + client.get_object.side_effect = _client_error(error_code) storage = ArchiveStorage(bucket=BUCKET_NAME) with pytest.raises(FileNotFoundError, match="Archive object not found"): storage.get_object("missing") +@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"]) +def test_get_object_non_missing_error_fails_closed(monkeypatch: pytest.MonkeyPatch, error_code: str): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.get_object.side_effect = _client_error(error_code) + storage = ArchiveStorage(bucket=BUCKET_NAME) + + with pytest.raises(ArchiveStorageError, match="Failed to download object"): + storage.get_object("key") + + +def test_get_object_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.get_object.side_effect = _network_error() + storage = ArchiveStorage(bucket=BUCKET_NAME) + + with pytest.raises(ArchiveStorageError, match="Failed to download object"): + storage.get_object("key") + + def test_get_object_stream(monkeypatch: pytest.MonkeyPatch): _configure_storage(monkeypatch) client, _ = _mock_client(monkeypatch) @@ -174,30 +200,80 @@ def test_get_object_stream(monkeypatch: pytest.MonkeyPatch): assert list(storage.get_object_stream("key")) == [b"a", b"b"] -def test_get_object_stream_missing(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"]) +def test_get_object_stream_missing(monkeypatch: pytest.MonkeyPatch, error_code: str): _configure_storage(monkeypatch) client, _ = _mock_client(monkeypatch) - client.get_object.side_effect = _client_error("NoSuchKey") + client.get_object.side_effect = _client_error(error_code) storage = ArchiveStorage(bucket=BUCKET_NAME) with pytest.raises(FileNotFoundError, match="Archive object not found"): list(storage.get_object_stream("missing")) -def test_object_exists(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"]) +def test_object_exists_returns_false_only_for_not_found( + monkeypatch: pytest.MonkeyPatch, + error_code: str, +): _configure_storage(monkeypatch) client, _ = _mock_client(monkeypatch) storage = ArchiveStorage(bucket=BUCKET_NAME) assert storage.object_exists("key") is True - client.head_object.side_effect = _client_error("404") + client.head_object.side_effect = _client_error(error_code) assert storage.object_exists("missing") is False -def test_delete_object_error(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"]) +def test_object_exists_raises_when_existence_is_unknown( + monkeypatch: pytest.MonkeyPatch, + error_code: str, +): _configure_storage(monkeypatch) client, _ = _mock_client(monkeypatch) - client.delete_object.side_effect = _client_error("500") + client.head_object.side_effect = _client_error(error_code) + storage = ArchiveStorage(bucket=BUCKET_NAME) + + with pytest.raises(ArchiveStorageError, match="Failed to check archive object"): + storage.object_exists("key") + + +def test_object_exists_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.head_object.side_effect = _network_error() + storage = ArchiveStorage(bucket=BUCKET_NAME) + + with pytest.raises(ArchiveStorageError, match="Failed to check archive object"): + storage.object_exists("key") + + +@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"]) +def test_delete_object_error(monkeypatch: pytest.MonkeyPatch, error_code: str): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.delete_object.side_effect = _client_error(error_code) + storage = ArchiveStorage(bucket=BUCKET_NAME) + + with pytest.raises(ArchiveStorageError, match="Failed to delete object"): + storage.delete_object("key") + + +@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"]) +def test_delete_object_missing_is_idempotent(monkeypatch: pytest.MonkeyPatch, error_code: str): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.delete_object.side_effect = _client_error(error_code) + storage = ArchiveStorage(bucket=BUCKET_NAME) + + storage.delete_object("missing") + + +def test_delete_object_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch): + _configure_storage(monkeypatch) + client, _ = _mock_client(monkeypatch) + client.delete_object.side_effect = _network_error() storage = ArchiveStorage(bucket=BUCKET_NAME) with pytest.raises(ArchiveStorageError, match="Failed to delete object"): diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py b/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py new file mode 100644 index 00000000000..5fefb149f3a --- /dev/null +++ b/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py @@ -0,0 +1,766 @@ +import json +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import MagicMock, call, patch + +import pytest + +from services.retention.workflow_run.bundle_archive_maintenance import ( + ARCHIVED_TABLES, + ArchiveBundleCatalogEntry, + BundleManifest, + BundleOperationResult, + BundleReference, + WorkflowRunBundleArchiveMaintenance, +) +from services.retention.workflow_run.constants import ( + ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME, + ARCHIVE_BUNDLE_DELETED_MARKER_NAME, + ARCHIVE_BUNDLE_FORMAT, + ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME, + ARCHIVE_BUNDLE_RESTORED_MARKER_NAME, + ARCHIVE_BUNDLE_SCHEMA_VERSION, +) + +TENANT_ID = "1251fe32-c0c7-4fe2-a7bd-a8105267faf5" +CATALOG_ID = "019f63b7-5ca4-7681-9ce0-800283608f39" +BUNDLE_ID = "bundle-a" + + +def _table_records( + **overrides: list[dict[str, Any]], +) -> dict[str, list[dict[str, Any]]]: + records = {table_name: [] for table_name in ARCHIVED_TABLES} + records.update(overrides) + return records + + +def _catalog_entry(*, catalog_id: str = CATALOG_ID, shard: str = "00-of-01") -> ArchiveBundleCatalogEntry: + return ArchiveBundleCatalogEntry( + catalog_id=catalog_id, + tenant_id=TENANT_ID, + year=2025, + month=3, + shard=shard, + bundle_id=BUNDLE_ID, + workflow_run_count=0, + row_count=0, + archive_bytes=0, + ) + + +def _manifest( + entry: ArchiveBundleCatalogEntry, + *, + bundle_id: str = BUNDLE_ID, + table_records: dict[str, list[dict[str, Any]]] | None = None, +) -> bytes: + object_prefix = WorkflowRunBundleArchiveMaintenance._catalog_object_prefix(entry) + records = table_records or _table_records() + tables = { + table_name: { + "row_count": len(records[table_name]), + "checksum": "", + "size_bytes": 0, + "object_key": f"{object_prefix}/{table_name}.parquet", + } + for table_name in ( + "workflow_runs", + "workflow_app_logs", + "workflow_node_executions", + "workflow_node_execution_offload", + "workflow_pauses", + "workflow_pause_reasons", + "workflow_trigger_logs", + ) + } + return json.dumps( + { + "schema_version": ARCHIVE_BUNDLE_SCHEMA_VERSION, + "archive_format": ARCHIVE_BUNDLE_FORMAT, + "tenant_id": entry.tenant_id, + "tenant_prefix": entry.tenant_id[0], + "year": entry.year, + "month": entry.month, + "shard": entry.shard, + "bundle_id": bundle_id, + "object_prefix": object_prefix, + "workflow_run_count": len(records["workflow_runs"]), + "workflow_node_execution_count": len(records["workflow_node_executions"]), + "tables": tables, + "run_ids": [str(record["id"]) for record in records["workflow_runs"]], + } + ).encode() + + +def _session_factory(session: MagicMock) -> MagicMock: + factory = MagicMock() + factory.return_value.__enter__.return_value = session + return factory + + +def _bundle_reference( + entry: ArchiveBundleCatalogEntry, + *, + table_records: dict[str, list[dict[str, Any]]] | None = None, +) -> BundleReference: + manifest = cast(BundleManifest, json.loads(_manifest(entry, table_records=table_records))) + return BundleReference( + catalog=entry, + object_prefix=manifest["object_prefix"], + manifest_key=f"{manifest['object_prefix']}/manifest.json", + manifest_size_bytes=0, + manifest=manifest, + ) + + +def _sample_archive_records() -> dict[str, list[dict[str, Any]]]: + return _table_records( + workflow_runs=[ + {"id": "run-1", "status": "succeeded"}, + {"id": "run-2", "status": "failed"}, + ], + workflow_app_logs=[ + {"id": "app-log-1", "workflow_run_id": "run-1"}, + ], + workflow_node_executions=[ + {"id": "node-1", "workflow_run_id": "run-1"}, + ], + workflow_node_execution_offload=[ + {"id": "offload-1", "node_execution_id": "node-1"}, + ], + workflow_pauses=[ + {"id": "pause-1", "workflow_run_id": "run-1"}, + ], + workflow_pause_reasons=[ + {"id": "reason-1", "pause_id": "pause-1"}, + ], + workflow_trigger_logs=[ + {"id": "trigger-1", "workflow_run_id": "run-1"}, + ], + ) + + +def test_catalog_discovery_is_ordered_and_limited_before_storage_io() -> None: + entry = _catalog_entry() + bundle = SimpleNamespace( + id=entry.catalog_id, + tenant_id=entry.tenant_id, + year=entry.year, + month=entry.month, + shard=entry.shard, + bundle_id=entry.bundle_id, + workflow_run_count=entry.workflow_run_count, + row_count=entry.row_count, + archive_bytes=entry.archive_bytes, + ) + session = MagicMock() + session.get.return_value = bundle + session.scalars.return_value = [bundle] + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, storage), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + entries = maintenance._list_catalog_entries( + tenant_ids=[TENANT_ID], + target_year=2025, + target_month=3, + after_catalog_id=CATALOG_ID, + limit=2, + ) + + statement = session.scalars.call_args.args[0] + rendered = str(statement) + assert "workflow_run_archive_bundles.year" in rendered + assert "workflow_run_archive_bundles.month" in rendered + assert "workflow_run_archive_bundles.id >" in rendered + assert "ORDER BY workflow_run_archive_bundles.id ASC" in rendered + assert "LIMIT" in rendered + assert entries == [entry] + storage.list_objects.assert_not_called() + + +def test_catalog_discovery_filters_and_validates_the_requested_shard() -> None: + entry = _catalog_entry(shard="03-of-16") + bundle = SimpleNamespace( + id=entry.catalog_id, + tenant_id=entry.tenant_id, + year=entry.year, + month=entry.month, + shard=entry.shard, + bundle_id=entry.bundle_id, + workflow_run_count=entry.workflow_run_count, + row_count=entry.row_count, + archive_bytes=entry.archive_bytes, + ) + session = MagicMock() + session.get.return_value = bundle + session.scalars.return_value = [bundle] + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + entries = maintenance._list_catalog_entries( + tenant_ids=None, + target_year=2025, + target_month=3, + after_catalog_id=CATALOG_ID, + limit=2, + shard="03-of-16", + ) + + statement = session.scalars.call_args.args[0] + rendered = str(statement) + assert "workflow_run_archive_bundles.shard =" in rendered + assert entries == [entry] + + session.get.return_value = SimpleNamespace( + year=2025, + month=3, + tenant_id=TENANT_ID, + shard="04-of-16", + ) + with pytest.raises(ValueError, match="requested archive shard"): + maintenance._list_catalog_entries( + tenant_ids=None, + target_year=2025, + target_month=3, + after_catalog_id=CATALOG_ID, + limit=2, + shard="03-of-16", + ) + + +def test_catalog_shard_preflight_rejects_mixed_layout_before_delete() -> None: + session = MagicMock() + session.scalars.return_value = ["00-of-01"] + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with pytest.raises(ValueError, match=r"unexpected shards.*00-of-01"): + maintenance.validate_catalog_shards( + target_year=2025, + target_month=3, + shard_total=16, + ) + + statement = session.scalars.call_args.args[0] + rendered = str(statement) + assert "workflow_run_archive_bundles.year" in rendered + assert "workflow_run_archive_bundles.month" in rendered + assert "workflow_run_archive_bundles.shard NOT IN" in rendered + + +def test_catalog_shard_preflight_accepts_an_expected_subset() -> None: + session = MagicMock() + session.scalars.return_value = [] + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + maintenance.validate_catalog_shards( + target_year=2025, + target_month=3, + shard_total=16, + ) + + +def test_catalog_shard_preflight_uses_requested_tenant_scope() -> None: + session = MagicMock() + session.scalars.return_value = [] + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + maintenance.validate_catalog_shards( + target_year=2025, + target_month=3, + shard_total=16, + tenant_ids=[TENANT_ID], + ) + + statement = session.scalars.call_args.args[0] + assert "workflow_run_archive_bundles.tenant_id IN" in str(statement) + + +@pytest.mark.parametrize( + ("cursor_bundle", "tenant_ids", "error_message"), + [ + (None, None, "does not exist"), + ( + SimpleNamespace(year=2024, month=3, tenant_id=TENANT_ID), + None, + "requested archive month", + ), + ( + SimpleNamespace(year=2025, month=3, tenant_id="other-tenant"), + [TENANT_ID], + "requested tenant scope", + ), + ], +) +def test_catalog_discovery_rejects_cursor_outside_requested_scope( + cursor_bundle: SimpleNamespace | None, + tenant_ids: list[str] | None, + error_message: str, +) -> None: + session = MagicMock() + session.get.return_value = cursor_bundle + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with pytest.raises(ValueError, match=error_message): + maintenance._list_catalog_entries( + tenant_ids=tenant_ids, + target_year=2025, + target_month=3, + after_catalog_id=CATALOG_ID, + limit=1, + ) + + session.scalars.assert_not_called() + + +def test_catalog_manifest_identity_mismatch_fails_closed() -> None: + entry = _catalog_entry() + storage = MagicMock() + storage.get_object.return_value = _manifest(entry, bundle_id="other-bundle") + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, storage), + session_factory=cast(MagicMock, _session_factory(MagicMock())), + ) + + with pytest.raises(ValueError, match="identity does not match catalog"): + maintenance._build_bundle_reference(cast(MagicMock, storage), entry) + + +def test_bundle_maintenance_locks_the_existing_catalog_row() -> None: + entry = _catalog_entry() + session = MagicMock() + session.scalar.return_value = entry.catalog_id + + WorkflowRunBundleArchiveMaintenance._lock_catalog_entry(session, entry) + + statement = session.scalar.call_args.args[0] + rendered = str(statement) + assert "workflow_run_archive_bundles.id" in rendered + assert "workflow_run_archive_bundles.tenant_id" in rendered + assert "FOR UPDATE" in rendered + + +def test_failure_and_dry_run_do_not_return_a_persistable_cursor() -> None: + entry = _catalog_entry() + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, storage), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_list_catalog_entries", return_value=[entry]), + patch.object(maintenance, "_build_bundle_reference", side_effect=RuntimeError("manifest unavailable")), + ): + failed_summary = maintenance.delete_batch( + tenant_ids=None, + target_year=2025, + target_month=3, + after_catalog_id=None, + limit=1, + ) + + assert failed_summary.bundles_failed == 1 + assert failed_summary.next_catalog_id is None + assert failed_summary.preview_next_catalog_id is None + + dry_run = WorkflowRunBundleArchiveMaintenance( + dry_run=True, + storage=cast(MagicMock, storage), + session_factory=cast(MagicMock, _session_factory(session)), + ) + bundle_ref = BundleReference( + catalog=entry, + object_prefix="object-prefix", + manifest_key="manifest.json", + manifest_size_bytes=0, + manifest=cast(BundleManifest, {}), + ) + successful_result = BundleOperationResult( + catalog_id=entry.catalog_id, + bundle_id=entry.bundle_id, + tenant_id=entry.tenant_id, + object_prefix=bundle_ref.object_prefix, + success=True, + ) + with ( + patch.object(dry_run, "_list_catalog_entries", return_value=[entry]), + patch.object(dry_run, "_build_bundle_reference", return_value=bundle_ref), + patch.object(dry_run, "_delete_bundle", return_value=successful_result), + ): + dry_run_summary = dry_run.delete_batch( + tenant_ids=None, + target_year=2025, + target_month=3, + after_catalog_id=None, + limit=1, + ) + + assert dry_run_summary.next_catalog_id is None + assert dry_run_summary.preview_next_catalog_id == entry.catalog_id + + +def test_live_archive_subset_accepts_full_partial_and_absent_live_data() -> None: + archive_records = _sample_archive_records() + manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest + partial_records = _table_records( + workflow_node_execution_offload=archive_records["workflow_node_execution_offload"], + workflow_pause_reasons=archive_records["workflow_pause_reasons"], + ) + + for live_records in (archive_records, partial_records, _table_records()): + WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset( + manifest, + archive_records, + live_records, + ) + + +def test_live_archive_subset_rejects_extra_rows() -> None: + archive_records = _sample_archive_records() + manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest + live_records = _table_records( + workflow_app_logs=[ + {"id": "extra-app-log", "workflow_run_id": "run-1"}, + ] + ) + + with pytest.raises(ValueError, match="rows missing from archive for workflow_app_logs"): + WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset( + manifest, + archive_records, + live_records, + ) + + +def test_live_archive_subset_rejects_content_mismatch() -> None: + archive_records = _sample_archive_records() + manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest + live_records = _table_records( + workflow_runs=[ + {"id": "run-1", "status": "failed"}, + ] + ) + + with pytest.raises(ValueError, match="subset content checksum mismatch for workflow_runs"): + WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset( + manifest, + archive_records, + live_records, + ) + + +def test_live_bundle_scope_includes_archived_ids_and_indirect_children() -> None: + archive_records = _sample_archive_records() + manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest + session = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + def select_live_parent_ids(_session, model, _run_ids): + if model.__tablename__ == "workflow_node_executions": + return ["live-node"] + if model.__tablename__ == "workflow_pauses": + return ["live-pause"] + raise AssertionError(f"unexpected model: {model}") + + with ( + patch.object(maintenance, "_select_ids_by_run_ids", side_effect=select_live_parent_ids), + patch.object(maintenance, "_load_records_by_column", return_value=[]) as load_records, + ): + maintenance._load_live_bundle_records( + session, + manifest, + archive_records, + lock=True, + ) + + queries = [ + ( + call.args[1].__tablename__, + call.args[2].key, + set(call.args[3]), + call.kwargs["lock"], + ) + for call in load_records.call_args_list + ] + assert ("workflow_pause_reasons", "pause_id", {"pause-1", "live-pause"}, True) in queries + assert ("workflow_pause_reasons", "id", {"reason-1"}, True) in queries + assert ("workflow_node_execution_offload", "node_execution_id", {"node-1", "live-node"}, True) in queries + assert ("workflow_node_execution_offload", "id", {"offload-1"}, True) in queries + assert ("workflow_app_logs", "workflow_run_id", {"run-1", "run-2"}, True) in queries + assert ("workflow_app_logs", "id", {"app-log-1"}, True) in queries + + +def test_delete_bundle_accepts_matching_partial_rows_and_deletes_only_that_subset() -> None: + entry = _catalog_entry() + archive_records = _sample_archive_records() + bundle_ref = _bundle_reference(entry, table_records=archive_records) + partial_records = _table_records( + workflow_node_execution_offload=archive_records["workflow_node_execution_offload"], + workflow_pause_reasons=archive_records["workflow_pause_reasons"], + ) + expected_deleted_counts = {table_name: len(partial_records[table_name]) for table_name in ARCHIVED_TABLES} + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_is_restore_started", return_value=False), + patch.object(maintenance, "_is_deleted", return_value=False), + patch.object( + maintenance, + "_validate_archive_object", + return_value=(bundle_ref.manifest, archive_records, 123), + ), + patch.object( + maintenance, + "_load_live_bundle_records", + side_effect=[partial_records, _table_records()], + ), + patch.object( + maintenance, + "_delete_bundle_rows", + return_value=expected_deleted_counts, + ) as delete_bundle_rows, + patch.object(maintenance, "_put_marker"), + patch.object(maintenance, "_mark_deleted") as mark_deleted, + patch.object(maintenance, "_delete_marker"), + ): + result = maintenance._delete_bundle(session, storage, bundle_ref) + + assert result.success + delete_bundle_rows.assert_called_once_with(session, partial_records) + session.commit.assert_called_once_with() + session.rollback.assert_not_called() + mark_deleted.assert_called_once_with(storage, bundle_ref.object_prefix) + + +def test_delete_bundle_marks_an_already_absent_source_without_deleting_rows() -> None: + entry = _catalog_entry() + archive_records = _sample_archive_records() + bundle_ref = _bundle_reference(entry, table_records=archive_records) + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_is_restore_started", return_value=False), + patch.object(maintenance, "_is_deleted", return_value=False), + patch.object( + maintenance, + "_validate_archive_object", + return_value=(bundle_ref.manifest, archive_records, 123), + ), + patch.object(maintenance, "_load_live_bundle_records", return_value=_table_records()), + patch.object(maintenance, "_delete_bundle_rows") as delete_bundle_rows, + patch.object(maintenance, "_mark_deleted") as mark_deleted, + patch.object(maintenance, "_delete_marker") as delete_marker, + ): + result = maintenance._delete_bundle(session, storage, bundle_ref) + + assert result.success + delete_bundle_rows.assert_not_called() + session.commit.assert_not_called() + session.rollback.assert_not_called() + mark_deleted.assert_called_once_with(storage, bundle_ref.object_prefix) + assert delete_marker.call_count == 2 + + +def test_delete_bundle_with_deleted_marker_rejects_remaining_orphan_children() -> None: + entry = _catalog_entry() + archive_records = _sample_archive_records() + bundle_ref = _bundle_reference(entry, table_records=archive_records) + live_records = _table_records( + workflow_node_execution_offload=archive_records["workflow_node_execution_offload"], + ) + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_is_restore_started", return_value=False), + patch.object(maintenance, "_is_deleted", return_value=True), + patch.object( + maintenance, + "_validate_archive_object", + return_value=(bundle_ref.manifest, archive_records, 123), + ), + patch.object(maintenance, "_load_live_bundle_records", return_value=live_records), + patch.object(maintenance, "_delete_bundle_rows") as delete_bundle_rows, + ): + result = maintenance._delete_bundle(session, storage, bundle_ref) + + assert not result.success + assert "Live rows exist for bundle with deleted marker" in result.error + delete_bundle_rows.assert_not_called() + session.commit.assert_not_called() + session.rollback.assert_called_once_with() + + +def test_delete_bundle_rejects_an_in_progress_restore() -> None: + entry = _catalog_entry() + bundle_ref = _bundle_reference(entry) + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_is_restore_started", return_value=True), + patch.object(maintenance, "_validate_archive_object") as validate_archive, + ): + result = maintenance._delete_bundle(session, storage, bundle_ref) + + assert not result.success + assert "reconcile restore before delete" in result.error + validate_archive.assert_not_called() + session.commit.assert_not_called() + session.rollback.assert_called_once_with() + + +def test_delete_bundle_rows_use_only_verified_primary_keys() -> None: + live_records = _sample_archive_records() + session = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with patch.object( + maintenance, + "_delete_by_column", + side_effect=lambda _session, _model, _column, values: len(values), + ) as delete_by_column: + deleted_counts = maintenance._delete_bundle_rows(session, live_records) + + expected_table_order = [ + "workflow_pause_reasons", + "workflow_node_execution_offload", + "workflow_trigger_logs", + "workflow_app_logs", + "workflow_node_executions", + "workflow_pauses", + "workflow_runs", + ] + assert [call.args[1].__tablename__ for call in delete_by_column.call_args_list] == expected_table_order + assert all(call.args[2].key == "id" for call in delete_by_column.call_args_list) + assert deleted_counts == {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES} + + +def test_restore_does_not_skip_an_interrupted_delete_without_deleted_marker() -> None: + entry = _catalog_entry() + session = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + + with ( + patch.object(maintenance, "_is_deleted", return_value=False), + patch.object(maintenance, "_is_delete_started", return_value=True), + patch.object(maintenance, "_validate_live_counts") as validate_live_counts, + ): + result = maintenance._restore_bundle(session, MagicMock(), _bundle_reference(entry)) + + assert not result.success + assert "reconcile delete first" in result.error + validate_live_counts.assert_not_called() + session.commit.assert_not_called() + + +def test_restore_does_not_skip_missing_source_rows_without_deleted_marker() -> None: + entry = _catalog_entry() + session = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, MagicMock()), + session_factory=cast(MagicMock, _session_factory(session)), + ) + bundle_ref = _bundle_reference(entry) + + with ( + patch.object(maintenance, "_is_deleted", return_value=False), + patch.object(maintenance, "_is_delete_started", return_value=False), + patch.object( + maintenance, + "_validate_live_counts", + side_effect=ValueError("source rows are missing"), + ) as validate_live_counts, + ): + result = maintenance._restore_bundle(session, MagicMock(), bundle_ref) + + assert not result.success + assert "source rows are missing" in result.error + validate_live_counts.assert_called_once_with(session, bundle_ref.manifest, expected_present=True) + session.commit.assert_not_called() + + +def test_restore_reconciles_a_started_marker_after_the_source_commit() -> None: + entry = _catalog_entry() + session = MagicMock() + storage = MagicMock() + maintenance = WorkflowRunBundleArchiveMaintenance( + storage=cast(MagicMock, storage), + session_factory=cast(MagicMock, _session_factory(session)), + ) + bundle_ref = _bundle_reference(entry) + + with ( + patch.object(maintenance, "_is_deleted", return_value=False), + patch.object(maintenance, "_is_delete_started", return_value=False), + patch.object(maintenance, "_is_restore_started", return_value=True), + patch.object(maintenance, "_validate_live_counts") as validate_live_counts, + patch.object(maintenance, "_mark_restored") as mark_restored, + ): + result = maintenance._restore_bundle(session, storage, bundle_ref) + + assert result.success + validate_live_counts.assert_called_once_with(session, bundle_ref.manifest, expected_present=True) + mark_restored.assert_called_once_with(storage, bundle_ref.object_prefix) + session.commit.assert_not_called() + + +def test_mark_restored_clears_stale_delete_marker_before_releasing_restore_fence() -> None: + storage = MagicMock() + object_prefix = "bundle-prefix" + operations = MagicMock() + + with ( + patch.object(WorkflowRunBundleArchiveMaintenance, "_delete_marker") as delete_marker, + patch.object(WorkflowRunBundleArchiveMaintenance, "_put_marker") as put_marker, + ): + operations.attach_mock(delete_marker, "delete") + operations.attach_mock(put_marker, "put") + WorkflowRunBundleArchiveMaintenance._mark_restored(storage, object_prefix) + + assert operations.mock_calls == [ + call.delete(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME), + call.put(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME), + call.delete(storage, object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME), + call.delete(storage, object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME), + ]