From b3c7a706cd7883dd97df30f1c80ae48643bbd773 Mon Sep 17 00:00:00 2001 From: zhaohao1004 Date: Thu, 16 Jul 2026 15:05:24 +0800 Subject: [PATCH] fix: handle partially deleted workflow archive bundles --- .../bundle_archive_maintenance.py | 260 +++++++++--- .../test_bundle_archive_maintenance.py | 388 +++++++++++++++++- 2 files changed, 583 insertions(+), 65 deletions(-) diff --git a/api/services/retention/workflow_run/bundle_archive_maintenance.py b/api/services/retention/workflow_run/bundle_archive_maintenance.py index 5b3d3da37c3..3dac45509da 100644 --- a/api/services/retention/workflow_run/bundle_archive_maintenance.py +++ b/api/services/retention/workflow_run/bundle_archive_maintenance.py @@ -201,15 +201,30 @@ 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. + 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. @@ -472,42 +487,58 @@ class WorkflowRunBundleArchiveMaintenance: result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id) try: validation_start = time.time() - if self._is_deleted(storage, bundle_ref.object_prefix): - self._validate_live_counts(session, bundle_ref.manifest, expected_present=False) - result.validation_time = time.time() - validation_start - result.success = True - result.elapsed_time = time.time() - start_time - return result + 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 - ): - result.validation_time = time.time() - validation_start + 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._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.validation_time = time.time() - validation_start 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) + 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 + 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) @@ -530,13 +561,18 @@ class WorkflowRunBundleArchiveMaintenance: 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 @@ -598,6 +634,19 @@ class WorkflowRunBundleArchiveMaintenance: 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, @@ -700,6 +749,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, @@ -778,26 +931,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( @@ -869,11 +1009,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, @@ -918,8 +1053,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, @@ -927,23 +1070,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: @@ -974,6 +1117,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) @@ -981,10 +1128,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: 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 index ef842adea32..6f30ed9f570 100644 --- 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 @@ -1,24 +1,40 @@ import json from types import SimpleNamespace -from typing import cast -from unittest.mock import MagicMock, patch +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_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION +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) -> ArchiveBundleCatalogEntry: return ArchiveBundleCatalogEntry( catalog_id=catalog_id, @@ -33,11 +49,17 @@ def _catalog_entry(*, catalog_id: str = CATALOG_ID) -> ArchiveBundleCatalogEntry ) -def _manifest(entry: ArchiveBundleCatalogEntry, *, bundle_id: str = BUNDLE_ID) -> bytes: +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": 0, + "row_count": len(records[table_name]), "checksum": "", "size_bytes": 0, "object_key": f"{object_prefix}/{table_name}.parquet", @@ -63,10 +85,10 @@ def _manifest(entry: ArchiveBundleCatalogEntry, *, bundle_id: str = BUNDLE_ID) - "shard": entry.shard, "bundle_id": bundle_id, "object_prefix": object_prefix, - "workflow_run_count": 0, - "workflow_node_execution_count": 0, + "workflow_run_count": len(records["workflow_runs"]), + "workflow_node_execution_count": len(records["workflow_node_executions"]), "tables": tables, - "run_ids": [], + "run_ids": [str(record["id"]) for record in records["workflow_runs"]], } ).encode() @@ -77,8 +99,12 @@ def _session_factory(session: MagicMock) -> MagicMock: return factory -def _bundle_reference(entry: ArchiveBundleCatalogEntry) -> BundleReference: - manifest = cast(BundleManifest, json.loads(_manifest(entry))) +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"], @@ -88,6 +114,33 @@ def _bundle_reference(entry: ArchiveBundleCatalogEntry) -> BundleReference: ) +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( @@ -182,6 +235,20 @@ def test_catalog_manifest_identity_mismatch_fails_closed() -> None: 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() @@ -243,6 +310,261 @@ def test_failure_and_dry_run_do_not_return_a_persistable_cursor() -> 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() @@ -288,3 +610,49 @@ def test_restore_does_not_skip_missing_source_rows_without_deleted_marker() -> N 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), + ]