feat(retention): add V2 workflow run archive bundlesa (#37747)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tests/integration_tests/services/retention/test_workflow_run_archiver.py

@@ -1,17 +1,17 @@
 import datetime
-import io
 import json
 import uuid
-import zipfile
 from unittest.mock import MagicMock, patch
 
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 
 from services.retention.workflow_run.archive_paid_plan_workflow_run import (
     ArchiveSummary,
     WorkflowRunArchiver,
 )
-from services.retention.workflow_run.constants import ARCHIVE_SCHEMA_VERSION
+from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION
 
 
 class TestWorkflowRunArchiverInit:
@@ -39,6 +39,22 @@ class TestWorkflowRunArchiverInit:
         with pytest.raises(ValueError, match="workers must be at least 1"):
             WorkflowRunArchiver(workers=0)
 
+    def test_run_shard_index_without_total_raises(self):
+        with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"):
+            WorkflowRunArchiver(run_shard_index=0)
+
+    def test_run_shard_total_without_index_raises(self):
+        with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"):
+            WorkflowRunArchiver(run_shard_total=4)
+
+    def test_run_shard_total_above_supported_range_raises(self):
+        with pytest.raises(ValueError, match="run_shard_total must be between 1 and 16"):
+            WorkflowRunArchiver(run_shard_index=0, run_shard_total=17)
+
+    def test_run_shard_index_must_be_less_than_total(self):
+        with pytest.raises(ValueError, match="run_shard_index must be between 0 and run_shard_total - 1"):
+            WorkflowRunArchiver(run_shard_index=4, run_shard_total=4)
+
     def test_valid_init_defaults(self):
         archiver = WorkflowRunArchiver(days=30, batch_size=50)
         assert archiver.days == 30
@@ -55,29 +71,93 @@ class TestWorkflowRunArchiverInit:
         assert archiver.end_before is not None
         assert archiver.workers == 2
 
+    def test_delete_after_archive_is_not_supported_for_bundle_archive(self):
+        with pytest.raises(ValueError, match="delete_after_archive is not supported by bundle archive"):
+            WorkflowRunArchiver(delete_after_archive=True)
+
+    def test_get_runs_batch_passes_shard_options(self):
+        repo = MagicMock()
+        repo.get_runs_batch_by_time_range.return_value = []
+        archiver = WorkflowRunArchiver(
+            tenant_prefixes=["0", "a"],
+            run_shard_index=1,
+            run_shard_total=4,
+            workflow_run_repo=repo,
+        )
+
+        archiver._get_runs_batch(None)
+
+        repo.get_runs_batch_by_time_range.assert_called_once()
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] == ["0", "a"]
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_index"] == 1
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_total"] == 4
+
+    def test_get_runs_batch_prefers_planned_tenant_ids_over_prefix_filter(self):
+        repo = MagicMock()
+        repo.get_runs_batch_by_time_range.return_value = []
+        archiver = WorkflowRunArchiver(
+            tenant_ids=["0tenant"],
+            tenant_prefixes=["0"],
+            paid_tenant_ids=["0tenant"],
+            workflow_run_repo=repo,
+        )
+
+        archiver._get_runs_batch(None)
+
+        repo.get_runs_batch_by_time_range.assert_called_once()
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["0tenant"]
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] is None
+
+    def test_get_runs_batch_uses_current_tenant_scan_scope(self):
+        repo = MagicMock()
+        repo.get_runs_batch_by_time_range.return_value = []
+        archiver = WorkflowRunArchiver(
+            tenant_ids=["tenant-a", "tenant-b"],
+            workflow_run_repo=repo,
+        )
+
+        archiver._get_runs_batch(None, tenant_scope=["tenant-b"])
+
+        repo.get_runs_batch_by_time_range.assert_called_once()
+        assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["tenant-b"]
+
+    def test_start_message_includes_shard(self):
+        archiver = WorkflowRunArchiver(tenant_prefixes=["0"], run_shard_index=1, run_shard_total=4)
+
+        message = archiver._build_start_message()
+
+        assert "tenant_prefixes=0" in message
+        assert "run_shard=1/4" in message
+
+    def test_start_message_summarizes_large_planned_tenant_list(self):
+        tenant_ids = [f"tenant-{index}" for index in range(11)]
+        archiver = WorkflowRunArchiver(tenant_ids=tenant_ids, tenant_prefixes=["0"])
+
+        message = archiver._build_start_message()
+
+        assert "tenant_ids=11 planned tenants" in message
+        assert "tenant-10" not in message
+
 
 class TestBuildArchiveBundle:
-    def test_bundle_contains_manifest_and_all_tables(self):
+    def test_bundle_contains_manifest_and_all_table_objects(self):
         archiver = WorkflowRunArchiver(days=90)
+        run = MagicMock()
+        run.id = str(uuid.uuid4())
+        run.tenant_id = str(uuid.uuid4())
+        run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
+        identity = archiver._build_bundle_identity([run])
+        table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
 
-        manifest_data = json.dumps({"schema_version": ARCHIVE_SCHEMA_VERSION}).encode("utf-8")
-        table_payloads = dict.fromkeys(archiver.ARCHIVED_TABLES, b"")
+        table_stats, table_payloads, manifest_data = archiver._build_archive_payload(identity, [run], table_data)
+        manifest = json.loads(manifest_data)
 
-        bundle_bytes = archiver._build_archive_bundle(manifest_data, table_payloads)
-
-        with zipfile.ZipFile(io.BytesIO(bundle_bytes), "r") as zf:
-            names = set(zf.namelist())
-            assert "manifest.json" in names
-            for table in archiver.ARCHIVED_TABLES:
-                assert f"{table}.jsonl" in names, f"Missing {table}.jsonl in bundle"
-
-    def test_bundle_missing_table_payload_raises(self):
-        archiver = WorkflowRunArchiver(days=90)
-        manifest_data = b"{}"
-        incomplete_payloads = {archiver.ARCHIVED_TABLES[0]: b"data"}
-
-        with pytest.raises(ValueError, match="Missing archive payload"):
-            archiver._build_archive_bundle(manifest_data, incomplete_payloads)
+        assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION
+        assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT
+        assert manifest["object_prefix"] == identity.object_prefix
+        assert set(table_payloads) == set(archiver.ARCHIVED_TABLES)
+        assert {stat.table_name for stat in table_stats} == set(archiver.ARCHIVED_TABLES)
+        assert pq.read_table(pa.BufferReader(table_payloads["workflow_runs"])).num_rows == 1
 
 
 class TestGenerateManifest:
@@ -88,25 +168,39 @@ class TestGenerateManifest:
         run = MagicMock()
         run.id = str(uuid.uuid4())
         run.tenant_id = str(uuid.uuid4())
-        run.app_id = str(uuid.uuid4())
-        run.workflow_id = str(uuid.uuid4())
         run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
+        identity = archiver._build_bundle_identity([run])
 
         stats = [
-            TableStats(table_name="workflow_runs", row_count=1, checksum="abc123", size_bytes=512),
-            TableStats(table_name="workflow_app_logs", row_count=2, checksum="def456", size_bytes=1024),
+            TableStats(
+                table_name="workflow_runs",
+                row_count=1,
+                checksum="abc123",
+                size_bytes=512,
+                object_key="workflow_runs.parquet",
+            ),
+            TableStats(
+                table_name="workflow_node_executions",
+                row_count=2,
+                checksum="def456",
+                size_bytes=1024,
+                object_key="workflow_node_executions.parquet",
+            ),
         ]
 
-        manifest = archiver._generate_manifest(run, stats)
+        manifest = archiver._generate_manifest(identity, [run], stats)
 
-        assert manifest["schema_version"] == ARCHIVE_SCHEMA_VERSION
-        assert manifest["workflow_run_id"] == run.id
+        assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION
+        assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT
+        assert manifest["bundle_id"] == identity.bundle_id
         assert manifest["tenant_id"] == run.tenant_id
-        assert manifest["app_id"] == run.app_id
+        assert manifest["workflow_run_count"] == 1
+        assert manifest["workflow_node_execution_count"] == 2
+        assert manifest["run_ids"] == [run.id]
         assert "tables" in manifest
         assert manifest["tables"]["workflow_runs"]["row_count"] == 1
         assert manifest["tables"]["workflow_runs"]["checksum"] == "abc123"
-        assert manifest["tables"]["workflow_app_logs"]["row_count"] == 2
+        assert manifest["tables"]["workflow_node_executions"]["row_count"] == 2
 
 
 class TestFilterPaidTenants:
@@ -163,6 +257,19 @@ class TestFilterPaidTenants:
 
         assert result == set()
 
+    def test_planned_paid_tenants_skip_billing_lookup(self):
+        archiver = WorkflowRunArchiver(days=90, paid_tenant_ids=["t1", "t3"])
+
+        with (
+            patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg,
+            patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing,
+        ):
+            cfg.BILLING_ENABLED = True
+            result = archiver._filter_paid_tenants({"t1", "t2", "t3"})
+
+        billing.get_plan_bulk_with_cache.assert_not_called()
+        assert result == {"t1", "t3"}
+
 
 class TestDryRunArchive:
     @patch("services.retention.workflow_run.archive_paid_plan_workflow_run.get_archive_storage")
@@ -175,3 +282,81 @@ class TestDryRunArchive:
         mock_get_storage.assert_not_called()
         assert isinstance(summary, ArchiveSummary)
         assert summary.runs_failed == 0
+
+    def test_dry_run_estimates_table_and_object_sizes(self):
+        archiver = WorkflowRunArchiver(days=90, dry_run=True)
+        run = MagicMock()
+        run.id = "run-1"
+        run.tenant_id = "tenant-1"
+        run.app_id = "app-1"
+        run.workflow_id = "workflow-1"
+        run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
+        table_data = {
+            "workflow_runs": [{"id": "run-1", "tenant_id": "tenant-1"}],
+            "workflow_app_logs": [{"id": "log-1", "workflow_run_id": "run-1"}],
+        }
+
+        with patch.object(archiver, "_extract_bundle_data", return_value=table_data):
+            result = archiver._archive_bundle(MagicMock(), None, [run])
+
+        stats_by_table = {stat.table_name: stat for stat in result.tables}
+        assert result.success is True
+        assert result.object_size_bytes > 0
+        assert stats_by_table["workflow_runs"].row_count == 1
+        assert stats_by_table["workflow_runs"].size_bytes > 0
+        assert stats_by_table["workflow_app_logs"].row_count == 1
+        assert stats_by_table["workflow_app_logs"].size_bytes > 0
+        assert stats_by_table["workflow_node_executions"].row_count == 0
+        assert stats_by_table["workflow_node_executions"].size_bytes > 0
+
+    def test_summary_merges_dry_run_estimates(self):
+        summary = ArchiveSummary()
+        result = MagicMock()
+        result.object_size_bytes = 128
+        result.tables = [
+            MagicMock(table_name="workflow_runs", row_count=1, size_bytes=64),
+            MagicMock(table_name="workflow_app_logs", row_count=2, size_bytes=32),
+        ]
+
+        WorkflowRunArchiver._merge_result_stats(summary, result)
+
+        assert summary.total_object_size_bytes == 128
+        assert summary.table_stats["workflow_runs"].row_count == 1
+        assert summary.table_stats["workflow_runs"].size_bytes == 64
+        assert summary.table_stats["workflow_app_logs"].row_count == 2
+        assert summary.table_stats["workflow_app_logs"].size_bytes == 32
+
+
+class TestArchiveRunIdempotency:
+    def test_locked_bundle_is_skipped(self):
+        archiver = WorkflowRunArchiver(days=90)
+        run = MagicMock()
+        run.id = "run-1"
+        run.tenant_id = "tenant-1"
+        run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
+
+        with (
+            patch.object(archiver, "_lock_runs_for_archive", return_value=[]),
+        ):
+            storage = MagicMock()
+            storage.object_exists.return_value = False
+            result = archiver._archive_bundle(MagicMock(), storage, [run])
+
+        assert result.success is True
+        assert result.skipped is True
+        assert result.error == "one or more runs locked or deleted by another archiver"
+
+    def test_already_archived_bundle_is_skipped(self):
+        archiver = WorkflowRunArchiver(days=90)
+        run = MagicMock()
+        run.id = "run-1"
+        run.tenant_id = "tenant-1"
+        run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0)
+        storage = MagicMock()
+        storage.object_exists.return_value = True
+
+        result = archiver._archive_bundle(MagicMock(), storage, [run])
+
+        assert result.success is True
+        assert result.skipped is True
+        assert result.error == "bundle already archived"

api/tests/test_containers_integration_tests/services/test_delete_archived_workflow_run.py

@@ -2,17 +2,44 @@
 Testcontainers integration tests for archived workflow run deletion service.
 """
 
+import io
+import json
+import zipfile
 from datetime import UTC, datetime, timedelta
+from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from graphon.enums import WorkflowExecutionStatus
+from libs.archive_storage import ArchiveStorage
 from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
 from models.workflow import WorkflowArchiveLog, WorkflowRun
+from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME, ARCHIVE_SCHEMA_VERSION
 from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
 
+ARCHIVED_TABLES = [
+    "workflow_runs",
+    "workflow_app_logs",
+    "workflow_node_executions",
+    "workflow_node_execution_offload",
+    "workflow_pauses",
+    "workflow_pause_reasons",
+    "workflow_trigger_logs",
+]
+
+
+class FakeArchiveStorage:
+    def __init__(self, objects: dict[str, bytes]):
+        self.objects = objects
+
+    def object_exists(self, key: str) -> bool:
+        return key in self.objects
+
+    def get_object(self, key: str) -> bytes:
+        return self.objects[key]
+
 
 class TestArchivedWorkflowRunDeletion:
     def _create_workflow_run(
@@ -47,7 +74,7 @@ class TestArchivedWorkflowRunDeletion:
         db_session_with_containers.commit()
         return run
 
-    def _create_archive_log(self, db_session_with_containers: Session, *, run: WorkflowRun) -> None:
+    def _create_archive_log(self, db_session_with_containers: Session, *, run: WorkflowRun) -> WorkflowArchiveLog:
         archive_log = WorkflowArchiveLog(
             tenant_id=run.tenant_id,
             app_id=run.app_id,
@@ -72,6 +99,59 @@ class TestArchivedWorkflowRunDeletion:
         )
         db_session_with_containers.add(archive_log)
         db_session_with_containers.commit()
+        return archive_log
+
+    def _archive_key(self, run: WorkflowRun) -> str:
+        return (
+            f"{run.tenant_id}/app_id={run.app_id}/year={run.created_at.strftime('%Y')}/"
+            f"month={run.created_at.strftime('%m')}/workflow_run_id={run.id}/{ARCHIVE_BUNDLE_NAME}"
+        )
+
+    def _archive_bundle(self, run: WorkflowRun, *, workflow_run_rows: int = 1) -> bytes:
+        table_payloads: dict[str, bytes] = {}
+        table_counts = {
+            "workflow_runs": workflow_run_rows,
+            "workflow_app_logs": 0,
+            "workflow_node_executions": 0,
+            "workflow_node_execution_offload": 0,
+            "workflow_pauses": 0,
+            "workflow_pause_reasons": 0,
+            "workflow_trigger_logs": 0,
+        }
+        for table_name in ARCHIVED_TABLES:
+            records = [{"id": run.id}] if table_name == "workflow_runs" and workflow_run_rows else []
+            table_payloads[table_name] = ArchiveStorage.serialize_to_jsonl(records)
+
+        manifest = {
+            "schema_version": ARCHIVE_SCHEMA_VERSION,
+            "workflow_run_id": run.id,
+            "tenant_id": run.tenant_id,
+            "app_id": run.app_id,
+            "workflow_id": run.workflow_id,
+            "created_at": run.created_at.isoformat(),
+            "archived_at": datetime.now(UTC).isoformat(),
+            "tables": {
+                table_name: {
+                    "row_count": table_counts[table_name],
+                    "checksum": ArchiveStorage.compute_checksum(payload),
+                    "size_bytes": len(payload),
+                }
+                for table_name, payload in table_payloads.items()
+            },
+        }
+        buffer = io.BytesIO()
+        with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
+            archive.writestr("manifest.json", json.dumps(manifest).encode("utf-8"))
+            for table_name, payload in table_payloads.items():
+                archive.writestr(f"{table_name}.jsonl", payload)
+        return buffer.getvalue()
+
+    def _patch_storage(self, run: WorkflowRun):
+        storage = FakeArchiveStorage({self._archive_key(run): self._archive_bundle(run)})
+        return patch(
+            "services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
+            return_value=storage,
+        )
 
     def test_delete_by_run_id_returns_error_when_run_missing(self, db_session_with_containers: Session):
         deleter = ArchivedWorkflowRunDeletion()
@@ -109,13 +189,23 @@ class TestArchivedWorkflowRunDeletion:
         self._create_archive_log(db_session_with_containers, run=run2)
         run_ids = [run1.id, run2.id]
 
-        deleter = ArchivedWorkflowRunDeletion()
-        results = deleter.delete_batch(
-            tenant_ids=[tenant_id],
-            start_date=base_time - timedelta(minutes=1),
-            end_date=base_time + timedelta(minutes=1),
-            limit=2,
+        storage = FakeArchiveStorage(
+            {
+                self._archive_key(run1): self._archive_bundle(run1),
+                self._archive_key(run2): self._archive_bundle(run2),
+            }
         )
+        deleter = ArchivedWorkflowRunDeletion()
+        with patch(
+            "services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
+            return_value=storage,
+        ):
+            results = deleter.delete_batch(
+                tenant_ids=[tenant_id],
+                start_date=base_time - timedelta(minutes=1),
+                end_date=base_time + timedelta(minutes=1),
+                limit=2,
+            )
 
         assert len(results) == 2
         assert all(result.success for result in results)
@@ -133,9 +223,11 @@ class TestArchivedWorkflowRunDeletion:
             created_at=datetime.now(UTC),
         )
         run_id = run.id
+        archive_log = self._create_archive_log(db_session_with_containers, run=run)
         deleter = ArchivedWorkflowRunDeletion()
 
-        result = deleter._delete_run(run)
+        with self._patch_storage(run):
+            result = deleter._delete_run(run, archive_log)
 
         assert result.success is True
         assert result.deleted_counts["runs"] == 1
@@ -152,9 +244,11 @@ class TestArchivedWorkflowRunDeletion:
             created_at=datetime.now(UTC),
         )
         run_id = run.id
+        archive_log = self._create_archive_log(db_session_with_containers, run=run)
         deleter = ArchivedWorkflowRunDeletion(dry_run=True)
 
-        result = deleter._delete_run(run)
+        with self._patch_storage(run):
+            result = deleter._delete_run(run, archive_log)
 
         assert result.success is True
         assert result.run_id == run_id
@@ -164,22 +258,33 @@ class TestArchivedWorkflowRunDeletion:
 
     def test_delete_run_exception_returns_error(self, db_session_with_containers: Session):
         """Exception during deletion should return failure result."""
-        from unittest.mock import MagicMock, patch
-
         tenant_id = str(uuid4())
         run = self._create_workflow_run(
             db_session_with_containers,
             tenant_id=tenant_id,
             created_at=datetime.now(UTC),
         )
+        archive_log = self._create_archive_log(db_session_with_containers, run=run)
         deleter = ArchivedWorkflowRunDeletion(dry_run=False)
 
+        expected_counts = {
+            "runs": 1,
+            "node_executions": 0,
+            "offloads": 0,
+            "app_logs": 0,
+            "trigger_logs": 0,
+            "pauses": 0,
+            "pause_reasons": 0,
+        }
         with patch.object(deleter, "_get_workflow_run_repo") as mock_get_repo:
             mock_repo = MagicMock()
             mock_get_repo.return_value = mock_repo
+            mock_repo.get_archived_log_by_run_id.return_value = archive_log
+            mock_repo.count_runs_with_related.return_value = expected_counts
             mock_repo.delete_runs_with_related.side_effect = Exception("Database error")
 
-            result = deleter._delete_run(run)
+            with self._patch_storage(run):
+                result = deleter._delete_run(run, archive_log)
 
         assert result.success is False
         assert result.error == "Database error"
@@ -197,7 +302,8 @@ class TestArchivedWorkflowRunDeletion:
         run_id = run.id
 
         deleter = ArchivedWorkflowRunDeletion()
-        result = deleter.delete_by_run_id(run_id)
+        with self._patch_storage(run):
+            result = deleter.delete_by_run_id(run_id)
 
         assert result.success is True
         db_session_with_containers.expunge_all()
@@ -212,3 +318,48 @@ class TestArchivedWorkflowRunDeletion:
 
         assert repo1 is repo2
         assert deleter.workflow_run_repo is repo1
+
+    def test_delete_run_fails_when_archive_object_missing(self, db_session_with_containers: Session):
+        tenant_id = str(uuid4())
+        run = self._create_workflow_run(
+            db_session_with_containers,
+            tenant_id=tenant_id,
+            created_at=datetime.now(UTC),
+        )
+        archive_log = self._create_archive_log(db_session_with_containers, run=run)
+        deleter = ArchivedWorkflowRunDeletion()
+        storage = FakeArchiveStorage({})
+
+        with patch(
+            "services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
+            return_value=storage,
+        ):
+            result = deleter._delete_run(run, archive_log)
+
+        assert result.success is False
+        assert result.error == f"Archive bundle not found: {self._archive_key(run)}"
+        db_session_with_containers.expire_all()
+        assert db_session_with_containers.get(WorkflowRun, run.id) is not None
+
+    def test_delete_run_fails_when_manifest_count_differs_from_live_rows(self, db_session_with_containers: Session):
+        tenant_id = str(uuid4())
+        run = self._create_workflow_run(
+            db_session_with_containers,
+            tenant_id=tenant_id,
+            created_at=datetime.now(UTC),
+        )
+        archive_log = self._create_archive_log(db_session_with_containers, run=run)
+        bundle = self._archive_bundle(run, workflow_run_rows=0)
+        storage = FakeArchiveStorage({self._archive_key(run): bundle})
+        deleter = ArchivedWorkflowRunDeletion()
+
+        with patch(
+            "services.retention.workflow_run.delete_archived_workflow_run.get_archive_storage",
+            return_value=storage,
+        ):
+            result = deleter._delete_run(run, archive_log)
+
+        assert result.success is False
+        assert "Archive row count mismatch before delete" in str(result.error)
+        db_session_with_containers.expire_all()
+        assert db_session_with_containers.get(WorkflowRun, run.id) is not None

api/tests/unit_tests/services/retention/workflow_run/test_delete_archived_workflow_run.py

@@ -0,0 +1,96 @@
+import io
+import json
+import zipfile
+from datetime import UTC, datetime
+
+import pytest
+
+from libs.archive_storage import ArchiveStorage
+from services.retention.workflow_run.constants import ARCHIVE_SCHEMA_VERSION
+from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
+
+ARCHIVED_TABLES = [
+    "workflow_runs",
+    "workflow_app_logs",
+    "workflow_node_executions",
+    "workflow_node_execution_offload",
+    "workflow_pauses",
+    "workflow_pause_reasons",
+    "workflow_trigger_logs",
+]
+
+
+def _build_archive_bundle(
+    *,
+    run_id: str = "run-1",
+    tenant_id: str = "tenant-1",
+    app_id: str = "app-1",
+    workflow_id: str = "workflow-1",
+    corrupt_checksum_for: str | None = None,
+) -> bytes:
+    table_payloads: dict[str, bytes] = {}
+    for table_name in ARCHIVED_TABLES:
+        records = [{"id": run_id}] if table_name == "workflow_runs" else []
+        table_payloads[table_name] = ArchiveStorage.serialize_to_jsonl(records)
+
+    manifest = {
+        "schema_version": ARCHIVE_SCHEMA_VERSION,
+        "workflow_run_id": run_id,
+        "tenant_id": tenant_id,
+        "app_id": app_id,
+        "workflow_id": workflow_id,
+        "created_at": datetime.now(UTC).isoformat(),
+        "archived_at": datetime.now(UTC).isoformat(),
+        "tables": {
+            table_name: {
+                "row_count": 1 if table_name == "workflow_runs" else 0,
+                "checksum": ArchiveStorage.compute_checksum(payload),
+                "size_bytes": len(payload),
+            }
+            for table_name, payload in table_payloads.items()
+        },
+    }
+    if corrupt_checksum_for:
+        manifest["tables"][corrupt_checksum_for]["checksum"] = "bad-checksum"
+
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
+        archive.writestr("manifest.json", json.dumps(manifest).encode("utf-8"))
+        for table_name, payload in table_payloads.items():
+            archive.writestr(f"{table_name}.jsonl", payload)
+    return buffer.getvalue()
+
+
+def test_validate_archive_bundle_accepts_valid_archive() -> None:
+    manifest = ArchivedWorkflowRunDeletion._validate_archive_bundle(
+        _build_archive_bundle(),
+        run_id="run-1",
+        tenant_id="tenant-1",
+        app_id="app-1",
+        workflow_id="workflow-1",
+    )
+
+    assert manifest["schema_version"] == ARCHIVE_SCHEMA_VERSION
+    assert manifest["tables"]["workflow_runs"]["row_count"] == 1
+
+
+def test_validate_archive_bundle_rejects_checksum_mismatch() -> None:
+    with pytest.raises(ValueError, match="archive member checksum mismatch"):
+        ArchivedWorkflowRunDeletion._validate_archive_bundle(
+            _build_archive_bundle(corrupt_checksum_for="workflow_runs"),
+            run_id="run-1",
+            tenant_id="tenant-1",
+            app_id="app-1",
+            workflow_id="workflow-1",
+        )
+
+
+def test_validate_archive_bundle_rejects_manifest_target_mismatch() -> None:
+    with pytest.raises(ValueError, match="manifest tenant_id does not match delete target"):
+        ArchivedWorkflowRunDeletion._validate_archive_bundle(
+            _build_archive_bundle(),
+            run_id="run-1",
+            tenant_id="different-tenant",
+            app_id="app-1",
+            workflow_id="workflow-1",
+        )

api/tests/unit_tests/services/test_archive_workflow_run_logs.py

@@ -9,8 +9,6 @@ This module contains tests for:
 from datetime import datetime
 from unittest.mock import MagicMock, patch
 
-from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_NAME
-
 
 class TestWorkflowRunArchiver:
     """Tests for the WorkflowRunArchiver class."""
@@ -37,18 +35,20 @@ class TestWorkflowRunArchiver:
         assert archiver.limit == 50
         assert archiver.dry_run is True
 
-    def test_get_archive_key(self):
-        """Test archive key generation."""
+    def test_get_bundle_manifest_key(self):
+        """Test V2 bundle manifest key generation."""
         from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver
 
-        archiver = WorkflowRunArchiver.__new__(WorkflowRunArchiver)
+        archiver = WorkflowRunArchiver(run_shard_index=1, run_shard_total=4)
 
         mock_run = MagicMock()
-        mock_run.tenant_id = "tenant-123"
-        mock_run.app_id = "app-999"
+        mock_run.tenant_id = "9enant-123"
         mock_run.id = "run-456"
         mock_run.created_at = datetime(2024, 1, 15, 12, 0, 0)
 
-        key = archiver._get_archive_key(mock_run)
+        identity = archiver._build_bundle_identity([mock_run])
+        key = archiver._get_manifest_object_key(identity)
 
-        assert key == f"tenant-123/app_id=app-999/year=2024/month=01/workflow_run_id=run-456/{ARCHIVE_BUNDLE_NAME}"
+        assert key.endswith("/manifest.json")
+        assert "workflow-runs/v2/tenant_prefix=9/tenant_id=9enant-123/year=2024/month=01" in key
+        assert "/shard=01-of-04/" in key