langgenius/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-19 16:44:00 -04:00
Code Issues 775 Packages Projects Releases 10 Wiki Activity

fix: handle partially deleted workflow archive bundles

Browse Source
This commit is contained in:
zhaohao1004
2026-07-16 15:05:24 +08:00
parent 9554037fec
commit b3c7a706cd
2 changed files with 583 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/services/retention/workflow_run/bundle_archive_maintenance.py
+205 -55
View File
@@ -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:
api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py
+378 -10
View File
@@ -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),
]