opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-21 10:38:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: use catalog for workflow archive maintenance (#39001)

Browse Source
This commit is contained in:
Ingram Z 2026-07-20 15:59:29 +08:00 committed by GitHub
parent f93a5b95c6
commit b7edc5ef0c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 2103 additions and 248 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

272
api/commands/retention.py
View File

@ -1,6 +1,8 @@
import datetime
import logging
import re
import time
import uuid
from collections.abc import Callable
from typing import TypedDict
@ -21,6 +23,7 @@ from tasks.remove_app_and_related_data_task import delete_draft_variables_batch
logger = logging.getLogger(__name__)
_HEX_PREFIXES = tuple("0123456789abcdef")
_TARGET_MONTH_PATTERN = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$")
class WorkflowRunArchivePlanRow(TypedDict):
@ -66,6 +69,7 @@ def _parse_tenant_prefixes(prefixes: str | None) -> list[str]:
def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[str] | None:
"""Keep an omitted scope unset while rejecting an explicitly empty scope."""
if raw_ids is None:
return None
parsed = sorted({raw_id.strip() for raw_id in raw_ids.split(",") if raw_id.strip()})
@ -74,6 +78,27 @@ def _parse_comma_separated_ids(raw_ids: str | None, *, param_name: str) -> list[
return parsed
def _parse_archive_target_month(target_month: str) -> tuple[int, int]:
"""Validate the V2 catalog month selector and return its numeric components."""
if not _TARGET_MONTH_PATTERN.fullmatch(target_month):
raise click.BadParameter("target-month must use YYYY-MM format", param_hint="--target-month")
year_text, month_text = target_month.split("-", maxsplit=1)
return int(year_text), int(month_text)
def _parse_archive_catalog_cursor(after_catalog_id: str | None) -> str | None:
"""Normalize the exclusive V2 catalog keyset cursor when one is provided."""
if after_catalog_id is None:
return None
try:
return str(uuid.UUID(after_catalog_id))
except ValueError as exc:
raise click.BadParameter(
"after-catalog-id must be a UUID returned by the same V2 operation and scope",
param_hint="--after-catalog-id",
) from exc
def _get_archive_candidate_tenant_ids_by_prefix(
session: Session,
prefix: str,
@ -810,9 +835,11 @@ def backfill_workflow_run_archive_bundles(
click.echo(click.style(f" ... and {len(summary.errors) - 10} more failures", fg="red"))
def _echo_bundle_archive_operation_summary(summary) -> None:
def _echo_bundle_archive_operation_summary(summary, *, dry_run: bool) -> None:
status = "completed successfully" if summary.bundles_failed == 0 else "completed with failures"
fg = "green" if summary.bundles_failed == 0 else "red"
cursor_label = "preview_next_catalog_id" if dry_run else "next_catalog_id"
cursor_value = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
click.echo(
click.style(
f"{summary.operation} {status}. "
@ -821,10 +848,12 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
f"archive_bytes={summary.archive_bytes} duration={summary.elapsed_time:.2f}s "
f"validation_time={summary.validation_time:.2f}s "
f"runs_per_second={summary.runs_per_second:.2f} rows_per_second={summary.rows_per_second:.2f} "
f"bytes_per_second={summary.bytes_per_second:.2f}",
f"bytes_per_second={summary.bytes_per_second:.2f} {cursor_label}={cursor_value or 'none'}",
fg=fg,
)
)
if dry_run:
click.echo(click.style("Dry-run cursor is preview-only; do not persist it for a destructive run.", fg="yellow"))
click.echo(click.style("table,row_count", fg="white"))
for table_name in [
"workflow_runs",
@ -842,7 +871,8 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
click.style(
f" bundle={result.bundle_id} tenant={result.tenant_id} runs={result.run_count} "
f"rows={result.row_count} archive_bytes={result.archive_bytes} "
f"time={result.elapsed_time:.2f}s validation={result.validation_time:.2f}s",
f"catalog_id={result.catalog_id} time={result.elapsed_time:.2f}s "
f"validation={result.validation_time:.2f}s",
fg="white",
)
)
@ -850,7 +880,7 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
click.echo(
click.style(
f" failed bundle={result.bundle_id} tenant={result.tenant_id} "
f"object_prefix={result.object_prefix} error={result.error}",
f"catalog_id={result.catalog_id} object_prefix={result.object_prefix} error={result.error}",
fg="red",
)
)
@ -867,25 +897,24 @@ def _echo_bundle_archive_operation_summary(summary) -> None:
)
@click.option("--run-id", required=False, help="Workflow run ID to restore.")
@click.option(
"--start-from",
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
"--target-month",
metavar="YYYY-MM",
default=None,
help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.",
help="V2 catalog month to restore; required unless --run-id is used.",
)
@click.option(
"--end-before",
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
"--after-catalog-id",
default=None,
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
help="Exclusive V2 cursor from the same restore month and tenant scope.",
)
@click.option("--workers", default=1, show_default=True, type=int, help="V1 --run-id compatibility only.")
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to restore.")
@click.option("--limit", type=click.IntRange(min=1), default=100, show_default=True, help="Maximum V2 catalog rows.")
@click.option("--dry-run", is_flag=True, help="Preview without restoring.")
def restore_workflow_runs(
tenant_ids: str | None,
run_id: str | None,
start_from: datetime.datetime | None,
end_before: datetime.datetime | None,
target_month: str | None,
after_catalog_id: str | None,
workers: int,
limit: int,
dry_run: bool,
@ -905,23 +934,20 @@ def restore_workflow_runs(
from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance
from services.retention.workflow_run.restore_archived_workflow_run import WorkflowRunRestore
parsed_tenant_ids = None
if tenant_ids:
parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()]
if not parsed_tenant_ids:
raise click.BadParameter("tenant-ids must not be empty")
parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids")
if (start_from is None) ^ (end_before is None):
raise click.UsageError("--start-from and --end-before must be provided together.")
if run_id is None and (start_from is None or end_before is None):
raise click.UsageError("--start-from and --end-before are required for batch restore.")
if workers < 1:
raise click.BadParameter("workers must be at least 1")
if run_id is not None and (target_month is not None or after_catalog_id is not None):
raise click.UsageError("--target-month and --after-catalog-id are only valid for V2 batch restore.")
if run_id is None and target_month is None:
raise click.UsageError("--target-month is required for V2 batch restore.")
start_time = datetime.datetime.now(datetime.UTC)
target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}"
click.echo(
click.style(
f"Starting restore of workflow run {run_id} at {start_time.isoformat()}.",
f"Starting restore of {target_desc} at {start_time.isoformat()}.",
fg="white",
)
)
@ -955,17 +981,20 @@ def restore_workflow_runs(
click.echo(
click.style("--workers is ignored for V2 bundle restore; bundles are processed serially.", fg="yellow")
)
assert start_from is not None
assert end_before is not None
assert target_month is not None
target_year, target_month_number = _parse_archive_target_month(target_month)
catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id)
bundle_restorer = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True)
summary = bundle_restorer.restore_batch(
tenant_ids=parsed_tenant_ids,
start_date=start_from,
end_date=end_before,
target_year=target_year,
target_month=target_month_number,
after_catalog_id=catalog_cursor,
limit=limit,
)
_echo_bundle_archive_operation_summary(summary)
return
_echo_bundle_archive_operation_summary(summary, dry_run=dry_run)
if summary.bundles_failed:
raise click.exceptions.Exit(1)
@click.command(
@ -979,23 +1008,41 @@ def restore_workflow_runs(
)
@click.option("--run-id", required=False, help="Workflow run ID to delete.")
@click.option(
"--start-from",
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
"--target-month",
metavar="YYYY-MM",
default=None,
help="Optional lower bound (inclusive) for created_at; must be paired with --end-before.",
help="V2 catalog month to delete; required unless --run-id is used.",
)
@click.option(
"--end-before",
type=click.DateTime(formats=["%Y-%m-%d", "%Y-%m-%dT%H:%M:%S"]),
"--after-catalog-id",
default=None,
help="Optional upper bound (exclusive) for created_at; must be paired with --start-from.",
help="Exclusive V2 cursor from the same delete month and tenant scope.",
)
@click.option(
"--run-shard-index",
default=None,
type=click.IntRange(min=0),
help="Zero-based archive shard index. Must be paired with --run-shard-total.",
)
@click.option(
"--run-shard-total",
default=None,
type=click.IntRange(min=1, max=16),
help="Total archive shard count. Must be paired with --run-shard-index.",
)
@click.option("--all-pages", is_flag=True, help="Process catalog pages until an empty page is reached.")
@click.option(
"--limit",
type=click.IntRange(min=1),
default=100,
show_default=True,
help="Maximum V2 catalog rows per page.",
)
@click.option("--limit", type=int, default=100, show_default=True, help="Maximum number of V2 bundles to delete.")
@click.option("--dry-run", is_flag=True, help="Preview without deleting.")
@click.option(
"--skip-bad-archives",
is_flag=True,
help="Continue batch deletion when one archive object fails validation.",
help="V1 --run-id only: continue when one archive object fails validation.",
)
@click.option(
"--restore-sample-interval",
@ -1007,8 +1054,11 @@ def restore_workflow_runs(
def delete_archived_workflow_runs(
tenant_ids: str | None,
run_id: str | None,
start_from: datetime.datetime | None,
end_before: datetime.datetime | None,
target_month: str | None,
after_catalog_id: str | None,
run_shard_index: int | None,
run_shard_total: int | None,
all_pages: bool,
limit: int,
dry_run: bool,
skip_bad_archives: bool,
@ -1018,26 +1068,38 @@ def delete_archived_workflow_runs(
Delete archived workflow runs from the database.
Batch delete uses V2 bundle metadata and validates object existence, manifest schema, object size, checksum, row
counts, and source/archive content checksums before deleting source rows. `--run-id` keeps the V1 per-run path.
counts, and source/archive content checksums before deleting source rows. Parallel workers may select one exact
archive shard; all-pages mode keeps only the current bounded page in memory. `--run-id` keeps the V1 per-run path.
"""
from services.retention.workflow_run.bundle_archive_maintenance import WorkflowRunBundleArchiveMaintenance
from services.retention.workflow_run.delete_archived_workflow_run import ArchivedWorkflowRunDeletion
parsed_tenant_ids = None
if tenant_ids:
parsed_tenant_ids = [tid.strip() for tid in tenant_ids.split(",") if tid.strip()]
if not parsed_tenant_ids:
raise click.BadParameter("tenant-ids must not be empty")
parsed_tenant_ids = _parse_comma_separated_ids(tenant_ids, param_name="tenant-ids")
if (start_from is None) ^ (end_before is None):
raise click.UsageError("--start-from and --end-before must be provided together.")
if run_id is None and (start_from is None or end_before is None):
raise click.UsageError("--start-from and --end-before are required for batch delete.")
if restore_sample_interval < 0:
raise click.BadParameter("restore-sample-interval must be >= 0")
if run_id is not None and (
target_month is not None
or after_catalog_id is not None
or run_shard_index is not None
or run_shard_total is not None
or all_pages
):
raise click.UsageError(
"--target-month, --after-catalog-id, --run-shard-index, --run-shard-total, and --all-pages "
"are only valid for V2 batch delete."
)
if run_id is None and target_month is None:
raise click.UsageError("--target-month is required for V2 batch delete.")
if run_id is None and skip_bad_archives:
raise click.UsageError("--skip-bad-archives is not supported for V2 catalog batches; they fail fast.")
if (run_shard_index is None) ^ (run_shard_total is None):
raise click.UsageError("--run-shard-index and --run-shard-total must be provided together.")
if run_shard_index is not None and run_shard_total is not None and run_shard_index >= run_shard_total:
raise click.UsageError("--run-shard-index must be less than --run-shard-total.")
start_time = datetime.datetime.now(datetime.UTC)
target_desc = f"workflow run {run_id}" if run_id else "workflow runs"
target_desc = f"workflow run {run_id}" if run_id else f"workflow archive catalog month {target_month}"
click.echo(
click.style(
f"Starting delete of {target_desc} at {start_time.isoformat()}.",
@ -1110,20 +1172,104 @@ def delete_archived_workflow_runs(
if restore_sample_interval:
click.echo(click.style("--restore-sample-interval is ignored for V2 bundle delete.", fg="yellow"))
assert start_from is not None
assert end_before is not None
bundle_deleter = WorkflowRunBundleArchiveMaintenance(
dry_run=dry_run,
strict_content_validation=True,
stop_on_error=not skip_bad_archives,
assert target_month is not None
target_year, target_month_number = _parse_archive_target_month(target_month)
catalog_cursor = _parse_archive_catalog_cursor(after_catalog_id)
shard = (
f"{run_shard_index:02d}-of-{run_shard_total:02d}"
if run_shard_index is not None and run_shard_total is not None
else None
)
summary = bundle_deleter.delete_batch(
tenant_ids=parsed_tenant_ids,
start_date=start_from,
end_date=end_before,
limit=limit,
)
_echo_bundle_archive_operation_summary(summary)
bundle_deleter = WorkflowRunBundleArchiveMaintenance(dry_run=dry_run, strict_content_validation=True)
if run_shard_total is not None:
try:
bundle_deleter.validate_catalog_shards(
target_year=target_year,
target_month=target_month_number,
shard_total=run_shard_total,
tenant_ids=parsed_tenant_ids,
)
except ValueError as exc:
logger.exception(
"Archive catalog shard preflight failed: target_month=%s shard=%s",
target_month,
shard,
)
raise click.ClickException(
f"Archive catalog shard preflight failed for target_month={target_month} shard={shard}: {exc}"
) from exc
initial_catalog_cursor = catalog_cursor
pages_processed = 0
bundles_succeeded = 0
runs_processed = 0
rows_processed = 0
archive_bytes = 0
while True:
summary = bundle_deleter.delete_batch(
tenant_ids=parsed_tenant_ids,
target_year=target_year,
target_month=target_month_number,
after_catalog_id=catalog_cursor,
limit=limit,
shard=shard,
)
_echo_bundle_archive_operation_summary(summary, dry_run=dry_run)
if summary.bundles_failed:
failed_result = next((result for result in summary.results if not result.success), None)
failed_catalog_id = failed_result.catalog_id if failed_result is not None else "unknown"
page_resume_cursor = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
resume_cursor = page_resume_cursor or catalog_cursor
if dry_run:
cursor_details = (
f"preview_after_catalog_id={resume_cursor or 'none'} "
f"destructive_retry_after_catalog_id={initial_catalog_cursor or 'none'}"
)
else:
cursor_details = f"resume_after_catalog_id={resume_cursor or 'none'}"
click.echo(
click.style(
f"Delete stopped: target_month={target_month} shard={shard or 'all'} "
f"failed_catalog_id={failed_catalog_id} "
f"{cursor_details}",
fg="red",
)
)
raise click.exceptions.Exit(1)
if not all_pages:
break
if summary.bundles_processed == 0:
break
pages_processed += 1
bundles_succeeded += summary.bundles_succeeded
runs_processed += summary.runs_processed
rows_processed += summary.rows_processed
archive_bytes += summary.archive_bytes
next_catalog_id = summary.preview_next_catalog_id if dry_run else summary.next_catalog_id
if next_catalog_id is None or (catalog_cursor is not None and next_catalog_id <= catalog_cursor):
click.echo(
click.style(
f"Delete cursor did not advance: target_month={target_month} shard={shard or 'all'} "
f"after_catalog_id={catalog_cursor or 'none'} next_catalog_id={next_catalog_id or 'none'}",
fg="red",
)
)
raise click.exceptions.Exit(1)
catalog_cursor = next_catalog_id
if all_pages:
final_cursor_label = "preview_final_catalog_id" if dry_run else "final_catalog_id"
click.echo(
click.style(
f"Delete all-pages completed successfully. target_month={target_month} shard={shard or 'all'} "
f"pages={pages_processed} bundles_success={bundles_succeeded} runs={runs_processed} "
f"rows={rows_processed} archive_bytes={archive_bytes} "
f"{final_cursor_label}={catalog_cursor or 'none'}",
fg="green",
)
)
def _find_orphaned_draft_variables(batch_size: int = 1000) -> list[str]:

45
api/libs/archive_storage.py
View File

@ -16,12 +16,19 @@ from urllib.parse import quote
import boto3
import orjson
from botocore.client import Config
from botocore.exceptions import ClientError
from botocore.exceptions import BotoCoreError, ClientError
from configs import dify_config
logger = logging.getLogger(__name__)
_OBJECT_NOT_FOUND_ERROR_CODES = frozenset({"404", "NoSuchKey", "NotFound"})
def _is_object_not_found_error(error: ClientError) -> bool:
error_code = str(error.response.get("Error", {}).get("Code", ""))
return error_code in _OBJECT_NOT_FOUND_ERROR_CODES
class ArchiveStorageError(Exception):
"""Base exception for archive storage operations."""
@ -138,10 +145,11 @@ class ArchiveStorage:
response = self.client.get_object(Bucket=self.bucket, Key=key)
return response["Body"].read()
except ClientError as e:
error_code = e.response.get("Error", {}).get("Code")
if error_code == "NoSuchKey":
raise FileNotFoundError(f"Archive object not found: {key}")
raise ArchiveStorageError(f"Failed to download object '{key}': {e}")
if _is_object_not_found_error(e):
raise FileNotFoundError(f"Archive object not found: {key}") from e
raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e
except BotoCoreError as e:
raise ArchiveStorageError(f"Failed to download object '{key}': {e}") from e
def get_object_stream(self, key: str) -> Generator[bytes, None, None]:
"""
@ -161,10 +169,11 @@ class ArchiveStorage:
response = self.client.get_object(Bucket=self.bucket, Key=key)
yield from response["Body"].iter_chunks()
except ClientError as e:
error_code = e.response.get("Error", {}).get("Code")
if error_code == "NoSuchKey":
raise FileNotFoundError(f"Archive object not found: {key}")
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}")
if _is_object_not_found_error(e):
raise FileNotFoundError(f"Archive object not found: {key}") from e
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e
except BotoCoreError as e:
raise ArchiveStorageError(f"Failed to stream object '{key}': {e}") from e
def object_exists(self, key: str) -> bool:
"""
@ -175,12 +184,19 @@ class ArchiveStorage:
Returns:
True if object exists, False otherwise
Raises:
ArchiveStorageError: If storage cannot authoritatively determine object existence
"""
try:
self.client.head_object(Bucket=self.bucket, Key=key)
return True
except ClientError:
return False
except ClientError as e:
if _is_object_not_found_error(e):
return False
raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e
except BotoCoreError as e:
raise ArchiveStorageError(f"Failed to check archive object '{key}': {e}") from e
def delete_object(self, key: str) -> None:
"""
@ -196,7 +212,12 @@ class ArchiveStorage:
self.client.delete_object(Bucket=self.bucket, Key=key)
logger.debug("Deleted object: %s", key)
except ClientError as e:
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}")
if _is_object_not_found_error(e):
logger.debug("Archive object was already absent: %s", key)
return
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e
except BotoCoreError as e:
raise ArchiveStorageError(f"Failed to delete object '{key}': {e}") from e
def generate_presigned_url(
self,

45
api/migrations/versions/2026_07_15_1600-3c9f8e2a1d7b_add_workflow_run_archive_bundle_cursor_indexes.py Normal file
View File

@ -0,0 +1,45 @@
"""add workflow-run archive bundle cursor indexes
Revision ID: 3c9f8e2a1d7b
Revises: 7a1c2d9e4b60
Create Date: 2026-07-15 16:00:00.000000
"""
from alembic import op
# revision identifiers, used by Alembic.
revision = "3c9f8e2a1d7b"
down_revision = "7a1c2d9e4b60"
branch_labels = None
depends_on = None
_TABLE_NAME = "workflow_run_archive_bundles"
_INDEXES = (
("workflow_run_archive_bundle_month_id_idx", ("year", "month", "id")),
("workflow_run_archive_bundle_month_shard_id_idx", ("year", "month", "shard", "id")),
)
def _is_postgresql() -> bool:
return op.get_bind().dialect.name == "postgresql"
def upgrade() -> None:
if _is_postgresql():
with op.get_context().autocommit_block():
for index_name, columns in _INDEXES:
op.create_index(index_name, _TABLE_NAME, columns, postgresql_concurrently=True)
return
for index_name, columns in _INDEXES:
op.create_index(index_name, _TABLE_NAME, columns)
def downgrade() -> None:
if _is_postgresql():
with op.get_context().autocommit_block():
for index_name, _ in reversed(_INDEXES):
op.drop_index(index_name, table_name=_TABLE_NAME, postgresql_concurrently=True)
return
for index_name, _ in reversed(_INDEXES):
op.drop_index(index_name, table_name=_TABLE_NAME)

2
api/models/workflow.py
View File

@ -1476,6 +1476,8 @@ class WorkflowRunArchiveBundle(DefaultFieldsDCMixin, TypeBase):
name="workflow_run_archive_bundle_identity_uq",
),
sa.Index("workflow_run_archive_bundle_tenant_month_idx", "tenant_id", "year", "month"),
sa.Index("workflow_run_archive_bundle_month_id_idx", "year", "month", "id"),
sa.Index("workflow_run_archive_bundle_month_shard_id_idx", "year", "month", "shard", "id"),
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)

41
api/services/retention/workflow_run/archive_paid_plan_workflow_run.py
View File

@ -6,7 +6,10 @@ This service archives workflow run logs for paid plan users older than the confi
Archive V2 writes bundle-level Parquet objects. A bundle contains many workflow runs and their related table rows.
Bundle metadata lives in the object-store manifest as the recoverable source of truth. Completed bundles are also
mirrored into a small database index so console listing and download jobs do not list object storage online.
published into a small database catalog so console listing, download, and maintenance jobs do not list object storage
online. Archive success requires that catalog publication to commit. A retry reconciles only its known manifest key
against an already-published shard index; a missing shard index fails closed rather than rebuilding it by scanning
historical manifests.
Archive campaigns should use fixed absolute UTC windows for every tenant-prefix/shard execution. Relative windows are
evaluated at process start and are not safe for multi-day rollout because each command would scan a different window.
@ -639,7 +642,6 @@ class WorkflowRunArchiver:
if storage is None:
raise ArchiveStorageNotConfiguredError("Archive storage not configured")
if storage.object_exists(self._get_manifest_object_key(identity)):
self._write_bundle_index(storage, identity)
self._sync_existing_bundle_index(session, storage, identity)
result.success = True
result.skipped = True
@ -651,6 +653,8 @@ class WorkflowRunArchiver:
runs = [run for run in runs if run.id not in archived_run_ids]
result.skipped_run_count = original_run_count - len(runs)
if not runs:
# Historical catalog rows are a rollout precondition. New bundles commit their catalog row before
# publishing the shard index, so this idempotency path never needs to rescan prior manifests.
result.run_count = 0
result.success = True
result.skipped = True
@ -663,7 +667,6 @@ class WorkflowRunArchiver:
result.object_prefix = identity.object_prefix
result.run_count = len(runs)
if storage.object_exists(self._get_manifest_object_key(identity)):
self._write_bundle_index(storage, identity)
self._sync_existing_bundle_index(session, storage, identity)
result.success = True
result.skipped = True
@ -695,10 +698,10 @@ class WorkflowRunArchiver:
for table_name, payload in table_payloads.items():
storage.put_object(self._get_table_object_key(identity, table_name), payload)
storage.put_object(self._get_manifest_object_key(identity), manifest_data)
self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs])
manifest = decode_archive_bundle_manifest(manifest_data)
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
session.commit()
self._merge_bundle_manifest_into_index(storage, identity, [run.id for run in runs])
logger.info(
"Archived workflow run bundle %s: tenant=%s runs=%s tables=%s object_prefix=%s",
@ -730,16 +733,18 @@ class WorkflowRunArchiver:
storage: ArchiveStorage,
identity: ArchiveBundleIdentity,
) -> None:
"""Best-effort DB index sync for a bundle whose manifest already exists in archive storage."""
"""Publish a known manifest to the DB catalog and reconcile its existing shard index."""
manifest_key = self._get_manifest_object_key(identity)
try:
manifest_data = storage.get_object(manifest_key)
manifest = decode_archive_bundle_manifest(manifest_data)
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
session.commit()
except Exception:
session.rollback()
logger.warning("Failed to sync workflow archive bundle index for %s", manifest_key, exc_info=True)
manifest_data = storage.get_object(manifest_key)
manifest = decode_archive_bundle_manifest(manifest_data)
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
session.commit()
self._merge_bundle_manifest_into_index(
storage,
identity,
manifest["run_ids"],
require_existing_index=True,
)
def _lock_runs_for_archive(
self,
@ -1029,10 +1034,20 @@ class WorkflowRunArchiver:
storage: ArchiveStorage,
identity: ArchiveBundleIdentity,
run_ids: Sequence[str],
*,
require_existing_index: bool = False,
) -> ArchiveBundleIndexDict:
"""
Merge one bundle into its shard index.
Retries for a known manifest set ``require_existing_index`` so they never create a partial index by scanning
or overwriting a shard whose historical entries cannot be proven from this one manifest.
"""
index_key = self._get_index_object_key(identity)
if storage.object_exists(index_key):
index = self._load_bundle_index(storage, identity)
elif require_existing_index:
raise RuntimeError(f"archive shard index missing while reconciling known manifest: {index_key}")
else:
index = self._build_bundle_index(storage, identity)

655
api/services/retention/workflow_run/bundle_archive_maintenance.py
View File

@ -1,14 +1,16 @@
"""
Maintain V2 workflow-run archive bundles.
Archive V2 keeps object-store manifests as the recoverable bundle source of truth. This maintenance module still
discovers delete/restore targets by listing `manifest.json` objects and uses object-store marker files for
delete/restore state. The separate database bundle index is intended for console listing and download jobs, not as the
source of truth for destructive maintenance.
Archive V2 keeps object-store manifests as the recoverable bundle source of truth. Delete and restore discover a
bounded page of candidates from `workflow_run_archive_bundles`, optionally restricted to one exact archive shard, then
construct each immutable manifest key from the catalog identity. They never list the object-store namespace.
Object-store marker files keep delete/restore idempotent, while the caller persists a non-dry-run cursor only after a
candidate succeeds.
Each bundle is processed in its own database transaction. A failed bundle leaves source rows unchanged unless the
transaction has already committed; marker handling makes the next run able to reconcile the common committed-but-marker
not-updated case.
not-updated case. Restore never skips a bundle with a missing deleted marker when deletion started or source rows have
drifted, so an external cursor cannot pass an interrupted delete.
"""
import datetime
@ -38,6 +40,7 @@ from models.workflow import (
WorkflowPause,
WorkflowPauseReason,
WorkflowRun,
WorkflowRunArchiveBundle,
)
from services.retention.workflow_run.constants import (
ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME,
@ -51,7 +54,6 @@ from services.retention.workflow_run.constants import (
logger = logging.getLogger(__name__)
_ARCHIVE_ROOT_PREFIX = "workflow-runs/v2/"
_CHUNK_SIZE = 5_000
@ -84,11 +86,28 @@ class BundleManifest(TypedDict):
@dataclass(frozen=True)
class BundleReference:
"""Object-store reference for one V2 archive bundle."""
class ArchiveBundleCatalogEntry:
"""Immutable catalog identity and manifest-derived metrics for one V2 archive bundle."""
catalog_id: str
tenant_id: str
year: int
month: int
shard: str
bundle_id: str
workflow_run_count: int
row_count: int
archive_bytes: int
@dataclass(frozen=True)
class BundleReference:
"""Verified object-store reference for one catalog candidate."""
catalog: ArchiveBundleCatalogEntry
object_prefix: str
manifest_key: str
manifest_size_bytes: int
manifest: BundleManifest
@ -96,6 +115,7 @@ class BundleReference:
class BundleOperationResult:
"""Result for one V2 bundle delete or restore operation."""
catalog_id: str
bundle_id: str
tenant_id: str
object_prefix: str
@ -128,6 +148,8 @@ class BundleOperationSummary:
archive_bytes: int = 0
elapsed_time: float = 0.0
validation_time: float = 0.0
next_catalog_id: str | None = None
preview_next_catalog_id: str | None = None
table_counts: dict[str, int] = field(default_factory=dict)
results: list[BundleOperationResult] = field(default_factory=list)
@ -180,162 +202,334 @@ RESTORE_ORDER = [
"workflow_trigger_logs",
]
DELETE_ORDER = [
"workflow_pause_reasons",
"workflow_node_execution_offload",
"workflow_trigger_logs",
"workflow_app_logs",
"workflow_node_executions",
"workflow_pauses",
"workflow_runs",
]
class WorkflowRunBundleArchiveMaintenance:
"""
Delete and restore V2 workflow-run archive bundles.
Delete accepts already-missing source rows only when every remaining row in the bundle scope is an unchanged
archive subset. It then removes only those verified primary keys, so unrelated or unarchived rows fail closed.
Non-dry-run delete and restore serialize on the existing archive catalog row before checking markers or changing
source rows.
Args:
dry_run: Validate and report counts without changing source rows or object-store markers.
strict_content_validation: Compare source-table content checksums against Parquet content before destructive
delete and after restore. Keep enabled for real maintenance.
stop_on_error: Stop batch processing after the first failed bundle.
strict_content_validation: Compare restored source-table content checksums against Parquet content. Delete
always validates that every remaining live row belongs to and matches the archive before removing it.
storage: Optional archive storage implementation. Tests may provide an in-memory implementation.
session_factory: Optional session factory. Each candidate is processed in its own transaction.
Batches stop at the first error so a returned cursor cannot pass an unhandled candidate.
"""
dry_run: bool
strict_content_validation: bool
stop_on_error: bool
storage: ArchiveStorage | None
session_factory: sessionmaker[Session]
def __init__(
self,
*,
dry_run: bool = False,
strict_content_validation: bool = True,
stop_on_error: bool = True,
storage: ArchiveStorage | None = None,
session_factory: sessionmaker[Session] | None = None,
) -> None:
self.dry_run = dry_run
self.strict_content_validation = strict_content_validation
self.stop_on_error = stop_on_error
self.storage = storage
self.session_factory = session_factory or sessionmaker(bind=db.engine, expire_on_commit=False)
def delete_batch(
self,
*,
tenant_ids: Sequence[str] | None,
start_date: datetime.datetime,
end_date: datetime.datetime,
target_year: int,
target_month: int,
after_catalog_id: str | None,
limit: int,
shard: str | None = None,
) -> BundleOperationSummary:
"""Validate and delete source rows for archived V2 bundles in the requested created_at window."""
"""Validate and delete one keyset page, optionally scoped to an exact archive shard."""
return self._process_batch(
operation="delete",
tenant_ids=tenant_ids,
start_date=start_date,
end_date=end_date,
target_year=target_year,
target_month=target_month,
after_catalog_id=after_catalog_id,
limit=limit,
shard=shard,
)
def restore_batch(
self,
*,
tenant_ids: Sequence[str] | None,
start_date: datetime.datetime,
end_date: datetime.datetime,
target_year: int,
target_month: int,
after_catalog_id: str | None,
limit: int,
) -> BundleOperationSummary:
"""Restore source rows for deleted V2 bundles in the requested created_at window."""
"""Restore source rows for a keyset page of deleted V2 bundles in one calendar month."""
return self._process_batch(
operation="restore",
tenant_ids=tenant_ids,
start_date=start_date,
end_date=end_date,
target_year=target_year,
target_month=target_month,
after_catalog_id=after_catalog_id,
limit=limit,
shard=None,
)
def validate_catalog_shards(
self,
*,
target_year: int,
target_month: int,
shard_total: int,
tenant_ids: Sequence[str] | None = None,
) -> None:
"""
Fail before a parallel delete when the requested closed-month scope contains a different shard layout.
A subset of the expected shards is valid because an archive shard may legitimately contain no bundles. Any
other shard name indicates a historical or mixed-layout month that must be handled by the serial delete path.
"""
if not 1 <= shard_total <= 16:
raise ValueError("shard_total must be between 1 and 16")
expected_shards = tuple(f"{index:02d}-of-{shard_total:02d}" for index in range(shard_total))
conditions = [
WorkflowRunArchiveBundle.year == target_year,
WorkflowRunArchiveBundle.month == target_month,
]
if tenant_ids is not None:
conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids))
statement = (
select(WorkflowRunArchiveBundle.shard)
.where(
*conditions,
WorkflowRunArchiveBundle.shard.not_in(expected_shards),
)
.distinct()
.order_by(WorkflowRunArchiveBundle.shard.asc())
)
with self.session_factory() as session:
unexpected_shards = list(session.scalars(statement))
if unexpected_shards:
raise ValueError(
"archive catalog month contains unexpected shards for "
f"{shard_total}-way delete: {', '.join(unexpected_shards)}"
)
def _process_batch(
self,
*,
operation: str,
tenant_ids: Sequence[str] | None,
start_date: datetime.datetime,
end_date: datetime.datetime,
target_year: int,
target_month: int,
after_catalog_id: str | None,
limit: int,
shard: str | None,
) -> BundleOperationSummary:
start_time = time.time()
summary = BundleOperationSummary(operation=operation)
if tenant_ids is not None and not tenant_ids:
return summary
storage = self._get_archive_storage()
bundle_refs = self._list_bundle_refs(
storage,
operation=operation,
storage = self.storage or self._get_archive_storage()
catalog_entries = self._list_catalog_entries(
tenant_ids=tenant_ids,
start_date=start_date,
end_date=end_date,
target_year=target_year,
target_month=target_month,
after_catalog_id=after_catalog_id,
limit=limit,
shard=shard,
)
logger.info("Found %s V2 archive bundles for %s", len(bundle_refs), operation)
session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
for bundle_ref in bundle_refs:
with session_maker() as session:
if operation == "delete":
result = self._delete_bundle(session, storage, bundle_ref)
elif operation == "restore":
result = self._restore_bundle(session, storage, bundle_ref)
else:
raise ValueError(f"Unsupported operation: {operation}")
logger.info(
"Found %s V2 archive catalog candidates for %s: year=%s month=%s shard=%s after_catalog_id=%s",
len(catalog_entries),
operation,
target_year,
target_month,
shard,
after_catalog_id,
)
for catalog_entry in catalog_entries:
try:
bundle_ref = self._build_bundle_reference(storage, catalog_entry)
with self.session_factory() as session:
if operation == "delete":
result = self._delete_bundle(session, storage, bundle_ref)
elif operation == "restore":
result = self._restore_bundle(session, storage, bundle_ref)
else:
raise ValueError(f"Unsupported operation: {operation}")
except Exception as exc:
result = self._new_result_from_catalog_entry(catalog_entry)
result.error = str(exc)
logger.exception(
"Failed to prepare V2 archive bundle %s from catalog %s",
catalog_entry.bundle_id,
catalog_entry.catalog_id,
)
self._merge_result(summary, result)
if not result.success and self.stop_on_error:
if result.success:
if self.dry_run:
summary.preview_next_catalog_id = catalog_entry.catalog_id
else:
summary.next_catalog_id = catalog_entry.catalog_id
else:
logger.error("Stopping V2 bundle %s after failure: %s", operation, result.error)
break
summary.elapsed_time = time.time() - start_time
return summary
def _list_bundle_refs(
def _list_catalog_entries(
self,
storage: ArchiveStorage,
*,
operation: str,
tenant_ids: Sequence[str] | None,
start_date: datetime.datetime,
end_date: datetime.datetime,
target_year: int,
target_month: int,
after_catalog_id: str | None,
limit: int,
) -> list[BundleReference]:
start_date = self._to_naive_utc(start_date)
end_date = self._to_naive_utc(end_date)
manifest_keys = self._list_manifest_keys(storage, tenant_ids)
refs: list[BundleReference] = []
for manifest_key in manifest_keys:
manifest_data = self._get_checked_object(storage, manifest_key)
object_prefix = manifest_key.removesuffix(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}")
manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix)
min_created_at = self._parse_manifest_datetime(manifest["min_created_at"])
max_created_at = self._parse_manifest_datetime(manifest["max_created_at"])
if max_created_at < start_date or min_created_at >= end_date:
continue
if tenant_ids and manifest["tenant_id"] not in tenant_ids:
continue
if operation == "delete" and self._is_deleted(storage, object_prefix):
continue
if operation == "restore" and not self._is_deleted(storage, object_prefix):
continue
refs.append(BundleReference(object_prefix=object_prefix, manifest_key=manifest_key, manifest=manifest))
shard: str | None = None,
) -> list[ArchiveBundleCatalogEntry]:
"""Read one bounded, stable-order candidate page from the database catalog."""
conditions = [
WorkflowRunArchiveBundle.year == target_year,
WorkflowRunArchiveBundle.month == target_month,
]
if tenant_ids:
conditions.append(WorkflowRunArchiveBundle.tenant_id.in_(tenant_ids))
if shard is not None:
conditions.append(WorkflowRunArchiveBundle.shard == shard)
if after_catalog_id:
conditions.append(WorkflowRunArchiveBundle.id > after_catalog_id)
refs.sort(
key=lambda ref: (
self._parse_manifest_datetime(ref.manifest["min_created_at"]),
ref.manifest["tenant_id"],
ref.manifest["bundle_id"],
)
statement = (
select(WorkflowRunArchiveBundle).where(*conditions).order_by(WorkflowRunArchiveBundle.id.asc()).limit(limit)
)
return refs[:limit]
with self.session_factory() as session:
if after_catalog_id:
cursor_bundle = session.get(WorkflowRunArchiveBundle, after_catalog_id)
self._validate_catalog_cursor_scope(
cursor_bundle,
tenant_ids=tenant_ids,
target_year=target_year,
target_month=target_month,
shard=shard,
)
bundles = list(session.scalars(statement))
return [
ArchiveBundleCatalogEntry(
catalog_id=bundle.id,
tenant_id=bundle.tenant_id,
year=bundle.year,
month=bundle.month,
shard=bundle.shard,
bundle_id=bundle.bundle_id,
workflow_run_count=bundle.workflow_run_count,
row_count=bundle.row_count,
archive_bytes=bundle.archive_bytes,
)
for bundle in bundles
]
@staticmethod
def _list_manifest_keys(storage: ArchiveStorage, tenant_ids: Sequence[str] | None) -> list[str]:
keys: list[str] = []
if tenant_ids:
prefixes = [
f"{_ARCHIVE_ROOT_PREFIX}tenant_prefix={tenant_id[0].lower()}/tenant_id={tenant_id}/"
for tenant_id in tenant_ids
]
else:
prefixes = [_ARCHIVE_ROOT_PREFIX]
for prefix in prefixes:
keys.extend(storage.list_objects(prefix))
return sorted(key for key in keys if key.endswith(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}"))
def _validate_catalog_cursor_scope(
cursor_bundle: WorkflowRunArchiveBundle | None,
*,
tenant_ids: Sequence[str] | None,
target_year: int,
target_month: int,
shard: str | None,
) -> None:
"""Reject a keyset cursor that cannot safely represent the requested catalog scope."""
if cursor_bundle is None:
raise ValueError("after_catalog_id does not exist in the workflow run archive bundle catalog")
if cursor_bundle.year != target_year or cursor_bundle.month != target_month:
raise ValueError("after_catalog_id is outside the requested archive month")
if tenant_ids is not None and cursor_bundle.tenant_id not in tenant_ids:
raise ValueError("after_catalog_id is outside the requested tenant scope")
if shard is not None and cursor_bundle.shard != shard:
raise ValueError("after_catalog_id is outside the requested archive shard")
def _build_bundle_reference(
self,
storage: ArchiveStorage,
catalog_entry: ArchiveBundleCatalogEntry,
) -> BundleReference:
"""Load and bind one manifest to the catalog row that selected it."""
object_prefix = self._catalog_object_prefix(catalog_entry)
manifest_key = f"{object_prefix}/{ARCHIVE_BUNDLE_MANIFEST_NAME}"
manifest_data = storage.get_object(manifest_key)
manifest = self._load_and_validate_manifest(manifest_data, object_prefix=object_prefix)
self._validate_manifest_catalog_identity(manifest, catalog_entry)
return BundleReference(
catalog=catalog_entry,
object_prefix=object_prefix,
manifest_key=manifest_key,
manifest_size_bytes=len(manifest_data),
manifest=manifest,
)
@staticmethod
def _catalog_object_prefix(catalog_entry: ArchiveBundleCatalogEntry) -> str:
"""Construct the immutable V2 bundle prefix from its database catalog identity."""
if not catalog_entry.tenant_id:
raise ValueError("archive catalog tenant_id must not be empty")
if not 1 <= catalog_entry.month <= 12:
raise ValueError(f"archive catalog month is invalid: {catalog_entry.month}")
return (
f"workflow-runs/v2/tenant_prefix={catalog_entry.tenant_id[0].lower()}/"
f"tenant_id={catalog_entry.tenant_id}/year={catalog_entry.year:04d}/"
f"month={catalog_entry.month:02d}/shard={catalog_entry.shard}/bundle={catalog_entry.bundle_id}"
)
@staticmethod
def _validate_manifest_catalog_identity(
manifest: BundleManifest,
catalog_entry: ArchiveBundleCatalogEntry,
) -> None:
"""Fail closed when the catalog locator and manifest identify different immutable bundles."""
expected_identity = (
catalog_entry.tenant_id,
catalog_entry.year,
catalog_entry.month,
catalog_entry.shard,
catalog_entry.bundle_id,
)
manifest_identity = (
manifest["tenant_id"],
manifest["year"],
manifest["month"],
manifest["shard"],
manifest["bundle_id"],
)
if manifest_identity != expected_identity:
raise ValueError(
f"archive manifest identity does not match catalog: expected={expected_identity}, "
f"actual={manifest_identity}"
)
if manifest["workflow_run_count"] != catalog_entry.workflow_run_count:
raise ValueError("archive manifest workflow_run_count does not match catalog")
manifest_row_count = sum(table["row_count"] for table in manifest["tables"].values())
if manifest_row_count != catalog_entry.row_count:
raise ValueError("archive manifest row_count does not match catalog")
def _delete_bundle(
self,
@ -344,37 +538,61 @@ class WorkflowRunBundleArchiveMaintenance:
bundle_ref: BundleReference,
) -> BundleOperationResult:
start_time = time.time()
result = self._new_result(bundle_ref.manifest)
result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id)
try:
validation_start = time.time()
if not self.dry_run:
self._lock_catalog_entry(session, bundle_ref.catalog)
if self._is_restore_started(storage, bundle_ref.object_prefix):
raise ValueError("restore started marker exists; reconcile restore before delete")
deleted_marker_exists = self._is_deleted(storage, bundle_ref.object_prefix)
manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref)
result.table_counts = self._manifest_table_counts(manifest)
result.archive_bytes = archive_bytes
self._lock_workflow_runs(session, manifest["run_ids"])
if self._is_delete_started(storage, bundle_ref.object_prefix) and self._live_counts_match(
session, manifest, expected_present=False
):
live_records = self._load_live_bundle_records(
session,
manifest,
table_records,
lock=not self.dry_run,
)
if deleted_marker_exists:
live_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES}
if any(live_counts.values()):
raise ValueError(f"Live rows exist for bundle with deleted marker: {live_counts}")
if not self.dry_run:
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
result.validation_time = time.time() - validation_start
result.success = True
result.elapsed_time = time.time() - start_time
return result
self._validate_live_archive_subset(manifest, table_records, live_records)
result.validation_time = time.time() - validation_start
if not any(live_records[table_name] for table_name in ARCHIVED_TABLES):
if not self.dry_run:
self._mark_deleted(storage, bundle_ref.object_prefix)
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
result.success = True
result.elapsed_time = time.time() - start_time
return result
self._validate_live_counts(session, manifest, expected_present=True)
if self.strict_content_validation:
self._validate_live_content(session, table_records)
result.validation_time = time.time() - validation_start
if not self.dry_run:
self._put_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
deleted_counts = self._delete_bundle_rows(session, table_records)
if deleted_counts != result.table_counts:
expected_deleted_counts = {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES}
deleted_counts = self._delete_bundle_rows(session, live_records)
if deleted_counts != expected_deleted_counts:
raise ValueError(
f"Deleted row count mismatch: expected={result.table_counts}, actual={deleted_counts}"
f"Deleted row count mismatch: expected={expected_deleted_counts}, actual={deleted_counts}"
)
self._validate_live_counts(session, manifest, expected_present=False)
remaining_records = self._load_live_bundle_records(session, manifest, table_records, lock=True)
remaining_counts = {table_name: len(remaining_records[table_name]) for table_name in ARCHIVED_TABLES}
if any(remaining_counts.values()):
raise ValueError(f"Live rows remain after bundle delete: {remaining_counts}")
session.commit()
self._mark_deleted(storage, bundle_ref.object_prefix)
self._delete_marker(storage, bundle_ref.object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME)
@ -394,9 +612,25 @@ class WorkflowRunBundleArchiveMaintenance:
bundle_ref: BundleReference,
) -> BundleOperationResult:
start_time = time.time()
result = self._new_result(bundle_ref.manifest)
result = self._new_result(bundle_ref.manifest, bundle_ref.catalog.catalog_id)
try:
validation_start = time.time()
if not self.dry_run:
self._lock_catalog_entry(session, bundle_ref.catalog)
if not self._is_deleted(storage, bundle_ref.object_prefix):
# A committed delete may be interrupted before `.deleted` is written. Do not let restore advance its
# cursor over that state: retry delete to reconcile it, or investigate source-row drift first.
if self._is_delete_started(storage, bundle_ref.object_prefix):
raise ValueError("delete started marker exists without a deleted marker; reconcile delete first")
restore_started = self._is_restore_started(storage, bundle_ref.object_prefix)
self._validate_live_counts(session, bundle_ref.manifest, expected_present=True)
result.validation_time = time.time() - validation_start
if restore_started and not self.dry_run:
self._mark_restored(storage, bundle_ref.object_prefix)
result.success = True
result.elapsed_time = time.time() - start_time
return result
manifest, table_records, archive_bytes = self._validate_archive_object(storage, bundle_ref)
result.table_counts = self._manifest_table_counts(manifest)
result.archive_bytes = archive_bytes
@ -408,6 +642,7 @@ class WorkflowRunBundleArchiveMaintenance:
if not self.dry_run:
self._mark_restored(storage, bundle_ref.object_prefix)
result.success = True
result.elapsed_time = time.time() - start_time
return result
self._validate_live_counts(session, manifest, expected_present=False)
@ -432,13 +667,40 @@ class WorkflowRunBundleArchiveMaintenance:
return result
@staticmethod
def _new_result(manifest: BundleManifest) -> BundleOperationResult:
def _new_result(manifest: BundleManifest, catalog_id: str) -> BundleOperationResult:
return BundleOperationResult(
catalog_id=catalog_id,
bundle_id=manifest["bundle_id"],
tenant_id=manifest["tenant_id"],
object_prefix=manifest["object_prefix"],
)
@staticmethod
def _new_result_from_catalog_entry(catalog_entry: ArchiveBundleCatalogEntry) -> BundleOperationResult:
try:
object_prefix = WorkflowRunBundleArchiveMaintenance._catalog_object_prefix(catalog_entry)
except ValueError:
object_prefix = "<invalid archive catalog identity>"
return BundleOperationResult(
catalog_id=catalog_entry.catalog_id,
bundle_id=catalog_entry.bundle_id,
tenant_id=catalog_entry.tenant_id,
object_prefix=object_prefix,
)
@staticmethod
def _lock_catalog_entry(session: Session, catalog_entry: ArchiveBundleCatalogEntry) -> None:
locked_catalog_id = session.scalar(
select(WorkflowRunArchiveBundle.id)
.where(
WorkflowRunArchiveBundle.id == catalog_entry.catalog_id,
WorkflowRunArchiveBundle.tenant_id == catalog_entry.tenant_id,
)
.with_for_update()
)
if locked_catalog_id is None:
raise ValueError("archive catalog row disappeared before bundle maintenance")
def _validate_archive_object(
self,
storage: ArchiveStorage,
@ -446,7 +708,7 @@ class WorkflowRunBundleArchiveMaintenance:
) -> tuple[BundleManifest, dict[str, list[dict[str, Any]]], int]:
manifest = bundle_ref.manifest
table_records: dict[str, list[dict[str, Any]]] = {}
total_size = len(storage.get_object(bundle_ref.manifest_key))
total_size = bundle_ref.manifest_size_bytes
for table_name in ARCHIVED_TABLES:
info = manifest["tables"][table_name]
payload = self._get_checked_object(storage, info["object_key"])
@ -469,12 +731,14 @@ class WorkflowRunBundleArchiveMaintenance:
f"expected={info['row_count']}, actual={len(records)}"
)
table_records[table_name] = records
if total_size != bundle_ref.catalog.archive_bytes:
raise ValueError(
f"Archive object total size mismatch: expected={bundle_ref.catalog.archive_bytes}, actual={total_size}"
)
return manifest, table_records, total_size
@staticmethod
def _get_checked_object(storage: ArchiveStorage, object_key: str) -> bytes:
if not storage.object_exists(object_key):
raise FileNotFoundError(f"Archive object not found: {object_key}")
return storage.get_object(object_key)
@staticmethod
@ -539,6 +803,110 @@ class WorkflowRunBundleArchiveMaintenance:
table = pq.read_table(io.BytesIO(payload))
return table.to_pylist()
def _load_live_bundle_records(
self,
session: Session,
manifest: BundleManifest,
table_records: dict[str, list[dict[str, Any]]],
*,
lock: bool,
) -> dict[str, list[dict[str, Any]]]:
"""Load the complete live scope for a bundle, including archived rows whose relationship fields drifted."""
run_ids = manifest["run_ids"]
archive_ids = {
table_name: [str(record["id"]) for record in table_records[table_name]] for table_name in ARCHIVED_TABLES
}
live_node_ids = self._select_ids_by_run_ids(session, WorkflowNodeExecutionModel, run_ids)
live_pause_ids = self._select_ids_by_run_ids(session, WorkflowPause, run_ids)
node_ids = sorted(set(archive_ids["workflow_node_executions"]) | set(live_node_ids))
pause_ids = sorted(set(archive_ids["workflow_pauses"]) | set(live_pause_ids))
def load_scope(
table_name: str,
model: Any,
scope_column: Any,
scope_ids: Sequence[str],
) -> list[dict[str, Any]]:
return self._merge_records_by_id(
self._load_records_by_column(session, model, scope_column, scope_ids, lock=lock),
self._load_records_by_column(session, model, model.id, archive_ids[table_name], lock=lock),
)
return {
"workflow_pause_reasons": load_scope(
"workflow_pause_reasons", WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
),
"workflow_node_execution_offload": load_scope(
"workflow_node_execution_offload",
WorkflowNodeExecutionOffload,
WorkflowNodeExecutionOffload.node_execution_id,
node_ids,
),
"workflow_trigger_logs": load_scope(
"workflow_trigger_logs", WorkflowTriggerLog, WorkflowTriggerLog.workflow_run_id, run_ids
),
"workflow_app_logs": load_scope(
"workflow_app_logs", WorkflowAppLog, WorkflowAppLog.workflow_run_id, run_ids
),
"workflow_node_executions": load_scope(
"workflow_node_executions",
WorkflowNodeExecutionModel,
WorkflowNodeExecutionModel.workflow_run_id,
run_ids,
),
"workflow_pauses": load_scope("workflow_pauses", WorkflowPause, WorkflowPause.workflow_run_id, run_ids),
"workflow_runs": self._load_records_by_column(
session,
WorkflowRun,
WorkflowRun.id,
sorted(set(run_ids) | set(archive_ids["workflow_runs"])),
lock=lock,
),
}
@classmethod
def _validate_live_archive_subset(
cls,
manifest: BundleManifest,
table_records: dict[str, list[dict[str, Any]]],
live_records: dict[str, list[dict[str, Any]]],
) -> None:
"""Require every live row in the bundle scope to exist unchanged in the validated archive."""
manifest_run_ids = {str(run_id) for run_id in manifest["run_ids"]}
if len(manifest_run_ids) != len(manifest["run_ids"]):
raise ValueError("archive manifest contains duplicate workflow run IDs")
archive_records_by_id: dict[str, dict[str, dict[str, Any]]] = {}
for table_name in ARCHIVED_TABLES:
records_by_id = {str(record["id"]): record for record in table_records[table_name]}
if len(records_by_id) != len(table_records[table_name]):
raise ValueError(f"archive contains duplicate row IDs for {table_name}")
archive_records_by_id[table_name] = records_by_id
if set(archive_records_by_id["workflow_runs"]) != manifest_run_ids:
raise ValueError("archive workflow run IDs do not match manifest run_ids")
for table_name in ARCHIVED_TABLES:
live_ids = [str(record["id"]) for record in live_records[table_name]]
if len(set(live_ids)) != len(live_ids):
raise ValueError(f"live scope contains duplicate row IDs for {table_name}")
archive_by_id = archive_records_by_id[table_name]
extra_ids = sorted(set(live_ids) - set(archive_by_id))
if extra_ids:
raise ValueError(
f"Live bundle scope contains rows missing from archive for {table_name}: {extra_ids[:10]}"
)
archive_subset = [archive_by_id[row_id] for row_id in live_ids]
live_checksum = cls._records_checksum(live_records[table_name])
archive_checksum = cls._records_checksum(archive_subset)
if live_checksum != archive_checksum:
raise ValueError(
f"Live/archive subset content checksum mismatch for {table_name}: "
f"expected={archive_checksum}, actual={live_checksum}"
)
def _validate_live_counts(
self,
session: Session,
@ -617,26 +985,13 @@ class WorkflowRunBundleArchiveMaintenance:
def _delete_bundle_rows(
self,
session: Session,
table_records: dict[str, list[dict[str, Any]]],
live_records: dict[str, list[dict[str, Any]]],
) -> dict[str, int]:
run_ids = [str(record["id"]) for record in table_records["workflow_runs"]]
node_ids = [str(record["id"]) for record in table_records["workflow_node_executions"]]
pause_ids = [str(record["id"]) for record in table_records["workflow_pauses"]]
deleted_counts = dict.fromkeys(ARCHIVED_TABLES, 0)
deleted_counts["workflow_pause_reasons"] = self._delete_by_column(
session, WorkflowPauseReason, WorkflowPauseReason.pause_id, pause_ids
)
deleted_counts["workflow_node_execution_offload"] = self._delete_by_column(
session, WorkflowNodeExecutionOffload, WorkflowNodeExecutionOffload.node_execution_id, node_ids
)
deleted_counts["workflow_trigger_logs"] = self._delete_by_run_ids(session, WorkflowTriggerLog, run_ids)
deleted_counts["workflow_app_logs"] = self._delete_by_run_ids(session, WorkflowAppLog, run_ids)
deleted_counts["workflow_node_executions"] = self._delete_by_run_ids(
session, WorkflowNodeExecutionModel, run_ids
)
deleted_counts["workflow_pauses"] = self._delete_by_run_ids(session, WorkflowPause, run_ids)
deleted_counts["workflow_runs"] = self._delete_by_run_ids(session, WorkflowRun, run_ids)
for table_name in DELETE_ORDER:
model = TABLE_MODELS[table_name]
row_ids = [str(record["id"]) for record in live_records[table_name]]
deleted_counts[table_name] = self._delete_by_column(session, model, model.id, row_ids)
return deleted_counts
def _restore_bundle_rows(
@ -708,11 +1063,6 @@ class WorkflowRunBundleArchiveMaintenance:
payload = json.dumps(normalized, sort_keys=True, default=str, ensure_ascii=False, separators=(",", ":"))
return ArchiveStorage.compute_checksum(payload.encode("utf-8"))
@staticmethod
def _lock_workflow_runs(session: Session, run_ids: Sequence[str]) -> None:
for chunk in WorkflowRunBundleArchiveMaintenance._chunks(run_ids, _CHUNK_SIZE):
list(session.scalars(select(WorkflowRun.id).where(WorkflowRun.id.in_(chunk)).with_for_update()))
@staticmethod
def _select_ids_by_run_ids(
session: Session,
@ -757,8 +1107,16 @@ class WorkflowRunBundleArchiveMaintenance:
session: Session,
model: Any,
run_ids: Sequence[str],
*,
lock: bool = False,
) -> list[dict[str, Any]]:
return self._load_records_by_column(session, model, self._run_id_column(model), run_ids)
return self._load_records_by_column(
session,
model,
self._run_id_column(model),
run_ids,
lock=lock,
)
def _load_records_by_column(
self,
@ -766,23 +1124,23 @@ class WorkflowRunBundleArchiveMaintenance:
model: Any,
column: Any,
values: Sequence[str],
*,
lock: bool = False,
) -> list[dict[str, Any]]:
if not values:
return []
rows: list[Any] = []
for chunk in self._chunks(values, _CHUNK_SIZE):
rows.extend(session.scalars(select(model).where(column.in_(chunk))))
statement = select(model).where(column.in_(chunk)).order_by(model.id.asc())
if lock:
statement = statement.with_for_update()
rows.extend(session.scalars(statement))
return [self._row_to_dict(row) for row in rows]
@staticmethod
def _delete_by_run_ids(
session: Session,
model: Any,
run_ids: Sequence[str],
) -> int:
return WorkflowRunBundleArchiveMaintenance._delete_by_column(
session, model, WorkflowRunBundleArchiveMaintenance._run_id_column(model), run_ids
)
def _merge_records_by_id(*record_groups: list[dict[str, Any]]) -> list[dict[str, Any]]:
records_by_id = {str(record["id"]): record for record_group in record_groups for record in record_group}
return [records_by_id[row_id] for row_id in sorted(records_by_id)]
@staticmethod
def _run_id_column(model: Any) -> Any:
@ -813,6 +1171,10 @@ class WorkflowRunBundleArchiveMaintenance:
def _is_delete_started(storage: ArchiveStorage, object_prefix: str) -> bool:
return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME}")
@staticmethod
def _is_restore_started(storage: ArchiveStorage, object_prefix: str) -> bool:
return storage.object_exists(f"{object_prefix}/{ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME}")
@staticmethod
def _mark_deleted(storage: ArchiveStorage, object_prefix: str) -> None:
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
@ -820,10 +1182,13 @@ class WorkflowRunBundleArchiveMaintenance:
@staticmethod
def _mark_restored(storage: ArchiveStorage, object_prefix: str) -> None:
WorkflowRunBundleArchiveMaintenance._delete_marker(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME)
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
WorkflowRunBundleArchiveMaintenance._delete_marker(
storage, object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME
)
WorkflowRunBundleArchiveMaintenance._delete_marker(
storage, object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME
)
WorkflowRunBundleArchiveMaintenance._put_marker(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME)
@staticmethod
def _put_marker(storage: ArchiveStorage, object_prefix: str, marker_name: str) -> None:
@ -836,16 +1201,6 @@ class WorkflowRunBundleArchiveMaintenance:
if storage.object_exists(marker_key):
storage.delete_object(marker_key)
@staticmethod
def _parse_manifest_datetime(value: str) -> datetime.datetime:
return WorkflowRunBundleArchiveMaintenance._to_naive_utc(datetime.datetime.fromisoformat(value))
@staticmethod
def _to_naive_utc(value: datetime.datetime) -> datetime.datetime:
if value.tzinfo is None:
return value
return value.astimezone(datetime.UTC).replace(tzinfo=None)
@staticmethod
def _chunks(values: Sequence[Any], size: int) -> list[Sequence[Any]]:
return [values[index : index + size] for index in range(0, len(values), size)]

107
api/tests/integration_tests/services/retention/test_workflow_run_archiver.py
View File

@ -513,11 +513,94 @@ class TestArchiveRunIdempotency:
storage = MagicMock()
storage.object_exists.return_value = True
result = archiver._archive_bundle(MagicMock(), storage, [run])
with patch.object(archiver, "_sync_existing_bundle_index") as sync_existing_bundle_index:
result = archiver._archive_bundle(MagicMock(), storage, [run])
assert result.success is True
assert result.skipped is True
assert result.error == "bundle already archived"
sync_existing_bundle_index.assert_called_once()
def test_existing_bundle_catalog_publication_failure_is_not_success(self):
archiver = WorkflowRunArchiver(days=90)
run = _run()
session = MagicMock()
storage = MagicMock()
storage.object_exists.return_value = True
with patch.object(archiver, "_sync_existing_bundle_index", side_effect=RuntimeError("catalog unavailable")):
result = archiver._archive_bundle(session, storage, [run])
assert result.success is False
assert result.error == "catalog unavailable"
session.rollback.assert_called_once()
def test_retry_repairs_index_after_catalog_commit_then_index_write_failure(self):
archiver = WorkflowRunArchiver(days=90)
run = _run()
identity = archiver._build_bundle_identity([run])
index_key = archiver._get_index_object_key(identity)
manifest_key = archiver._get_manifest_object_key(identity)
storage = FakeArchiveStorage()
original_put_object = storage.put_object
index_write_count = 0
def put_object(key: str, data: bytes) -> str:
nonlocal index_write_count
if key == index_key:
index_write_count += 1
if index_write_count == 2:
raise RuntimeError("index write failed")
return original_put_object(key, data)
storage.put_object = MagicMock(side_effect=put_object)
first_session = MagicMock()
first_session.scalar.return_value = None
table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
with (
patch.object(archiver, "_lock_runs_for_archive", return_value=[run]),
patch.object(archiver, "_extract_bundle_data", return_value=table_data),
):
first_result = archiver._archive_bundle(first_session, storage, [run])
assert first_result.success is False
assert first_result.error == "index write failed"
assert manifest_key in storage.objects
assert json.loads(storage.objects[index_key])["run_ids"] == []
storage.list_objects = MagicMock(wraps=storage.list_objects)
retry_session = MagicMock()
retry_session.scalar.return_value = None
retry_result = archiver._archive_bundle(retry_session, storage, [run])
assert retry_result.success is True
assert retry_result.skipped is True
assert json.loads(storage.objects[index_key])["manifest_keys"] == [manifest_key]
assert json.loads(storage.objects[index_key])["run_ids"] == [run.id]
storage.list_objects.assert_not_called()
def test_existing_manifest_with_missing_index_fails_without_partial_rebuild(self):
archiver = WorkflowRunArchiver(days=90)
run = _run()
identity = archiver._build_bundle_identity([run])
_, _, manifest_data = archiver._build_archive_payload(
identity,
[run],
{"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]},
)
manifest_key = archiver._get_manifest_object_key(identity)
index_key = archiver._get_index_object_key(identity)
storage = FakeArchiveStorage({manifest_key: manifest_data})
storage.list_objects = MagicMock(wraps=storage.list_objects)
result = archiver._archive_bundle(MagicMock(), storage, [run])
assert result.success is False
assert "archive shard index missing" in (result.error or "")
assert index_key not in storage.objects
storage.list_objects.assert_not_called()
def test_successful_bundle_persists_archive_index(self):
archiver = WorkflowRunArchiver(days=90)
@ -550,6 +633,28 @@ class TestArchiveRunIdempotency:
assert archived_bundle.row_count == 2
session.commit.assert_called_once()
def test_new_bundle_catalog_commit_failure_is_not_success(self):
archiver = WorkflowRunArchiver(days=90)
run = _run(str(uuid.uuid4()))
run.tenant_id = str(uuid.uuid4())
session = MagicMock()
session.scalar.return_value = None
session.commit.side_effect = RuntimeError("catalog commit failed")
storage = MagicMock()
storage.object_exists.return_value = False
storage.list_objects.return_value = []
table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}
with (
patch.object(archiver, "_lock_runs_for_archive", return_value=[run]),
patch.object(archiver, "_extract_bundle_data", return_value=table_data),
):
result = archiver._archive_bundle(session, storage, [run])
assert result.success is False
assert result.error == "catalog commit failed"
session.rollback.assert_called_once()
def test_index_skips_all_already_archived_runs(self):
archiver = WorkflowRunArchiver(days=90)
run = MagicMock()

324
api/tests/unit_tests/commands/test_archive_workflow_runs.py
View File

@ -3,9 +3,19 @@ from unittest.mock import MagicMock
import click
import pytest
from click.testing import CliRunner
from sqlalchemy.exc import OperationalError
from commands import retention
from services.retention.workflow_run import bundle_archive_maintenance
from services.retention.workflow_run.bundle_archive_maintenance import (
BundleOperationResult,
BundleOperationSummary,
)
_CURSOR_0 = "00000000-0000-0000-0000-000000000000"
_CURSOR_1 = "00000000-0000-0000-0000-000000000001"
_CURSOR_2 = "00000000-0000-0000-0000-000000000002"
def _db_disconnect_error() -> OperationalError:
@ -24,6 +34,58 @@ def _session_context(session):
return context
def _delete_summary(
*,
processed: int,
succeeded: int = 0,
failed: int = 0,
next_catalog_id: str | None = None,
preview_next_catalog_id: str | None = None,
results: list[BundleOperationResult] | None = None,
) -> BundleOperationSummary:
return BundleOperationSummary(
operation="delete",
bundles_processed=processed,
bundles_succeeded=succeeded,
bundles_failed=failed,
next_catalog_id=next_catalog_id,
preview_next_catalog_id=preview_next_catalog_id,
results=results or [],
)
def _patch_bundle_deleter(
monkeypatch: pytest.MonkeyPatch,
summaries: list[BundleOperationSummary],
) -> MagicMock:
deleter = MagicMock()
deleter.delete_batch.side_effect = summaries
monkeypatch.setattr(
bundle_archive_maintenance,
"WorkflowRunBundleArchiveMaintenance",
MagicMock(return_value=deleter),
)
return deleter
@pytest.mark.parametrize(
"command",
[retention.restore_workflow_runs, retention.delete_archived_workflow_runs],
)
def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command):
result = CliRunner().invoke(
command,
["--tenant-ids", "", "--target-month", "2025-03"],
)
assert result.exit_code == 2
assert "tenant-ids must not be empty" in result.output
def test_archive_tenant_id_parser_keeps_omitted_scope_unset():
assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None
def test_resolve_archive_tenant_ids_from_plan_uses_explicit_sessions(monkeypatch):
end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC)
sessions = [MagicMock(name="session-a"), MagicMock(name="session-b")]
@ -137,3 +199,265 @@ def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(mon
dry_run=True,
delete_after_archive=False,
)
def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages(monkeypatch):
deleter = _patch_bundle_deleter(
monkeypatch,
[_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--limit", "2"],
)
assert result.exit_code == 0
deleter.delete_batch.assert_called_once()
assert deleter.delete_batch.call_args.kwargs["target_year"] == 2025
assert deleter.delete_batch.call_args.kwargs["target_month"] == 3
assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] is None
assert deleter.delete_batch.call_args.kwargs["limit"] == 2
def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch):
deleter = _patch_bundle_deleter(
monkeypatch,
[
_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1),
_delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_2),
_delete_summary(processed=0),
],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--all-pages", "--limit", "2"],
)
assert result.exit_code == 0
assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [
None,
_CURSOR_1,
_CURSOR_2,
]
def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page(monkeypatch):
deleter = _patch_bundle_deleter(
monkeypatch,
[
_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1),
_delete_summary(processed=0),
],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--all-pages", "--limit", "2"],
)
assert result.exit_code == 0
assert deleter.delete_batch.call_count == 2
assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1
def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch):
failed_result = BundleOperationResult(
catalog_id=_CURSOR_2,
bundle_id="bundle-failed",
tenant_id="tenant-1",
object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed",
error="archive checksum mismatch",
)
deleter = _patch_bundle_deleter(
monkeypatch,
[
_delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_1),
_delete_summary(processed=1, failed=1, results=[failed_result]),
_delete_summary(processed=0),
],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--all-pages"],
)
assert result.exit_code == 1
assert deleter.delete_batch.call_count == 2
assert "target_month=2025-03" in result.output
assert f"failed_catalog_id={_CURSOR_2}" in result.output
assert f"resume_after_catalog_id={_CURSOR_1}" in result.output
def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance(monkeypatch):
deleter = _patch_bundle_deleter(
monkeypatch,
[_delete_summary(processed=1, succeeded=1, next_catalog_id=None)],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--all-pages"],
)
assert result.exit_code == 1
deleter.delete_batch.assert_called_once()
assert "cursor did not advance" in result.output.lower()
def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch):
deleter = _patch_bundle_deleter(
monkeypatch,
[
_delete_summary(processed=1, succeeded=1, preview_next_catalog_id=_CURSOR_1),
_delete_summary(processed=0),
],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", "--all-pages", "--dry-run"],
)
assert result.exit_code == 0
assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [
None,
_CURSOR_1,
]
def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors(monkeypatch):
failed_result = BundleOperationResult(
catalog_id=_CURSOR_2,
bundle_id="bundle-failed",
tenant_id="tenant-1",
object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed",
error="archive checksum mismatch",
)
deleter = _patch_bundle_deleter(
monkeypatch,
[
_delete_summary(
processed=1,
succeeded=1,
preview_next_catalog_id=_CURSOR_1,
),
_delete_summary(
processed=1,
failed=1,
results=[failed_result],
),
],
)
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
[
"--target-month",
"2025-03",
"--after-catalog-id",
_CURSOR_0,
"--all-pages",
"--dry-run",
],
)
assert result.exit_code == 1
assert deleter.delete_batch.call_count == 2
assert f"failed_catalog_id={_CURSOR_2}" in result.output
assert f"preview_after_catalog_id={_CURSOR_1}" in result.output
assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output
def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch):
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
[
"--target-month",
"2025-03",
"--after-catalog-id",
_CURSOR_0,
"--all-pages",
],
)
assert result.exit_code == 0
deleter.delete_batch.assert_called_once()
assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] == _CURSOR_0
@pytest.mark.parametrize(
"shard_args",
[
["--run-shard-index", "0"],
["--run-shard-total", "16"],
["--run-shard-index", "16", "--run-shard-total", "16"],
["--run-shard-index", "-1", "--run-shard-total", "16"],
["--run-shard-index", "0", "--run-shard-total", "0"],
["--run-shard-index", "0", "--run-shard-total", "17"],
],
)
def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args):
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
["--target-month", "2025-03", *shard_args],
)
assert result.exit_code == 2
deleter.delete_batch.assert_not_called()
def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch):
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
[
"--target-month",
"2025-03",
"--tenant-ids",
"tenant-1",
"--run-shard-index",
"3",
"--run-shard-total",
"16",
],
)
assert result.exit_code == 0
deleter.validate_catalog_shards.assert_called_once_with(
target_year=2025,
target_month=3,
shard_total=16,
tenant_ids=["tenant-1"],
)
deleter.delete_batch.assert_called_once()
assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16"
def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch):
deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01")
result = CliRunner().invoke(
retention.delete_archived_workflow_runs,
[
"--target-month",
"2025-03",
"--run-shard-index",
"3",
"--run-shard-total",
"16",
],
)
assert result.exit_code == 1
assert "shard preflight failed" in result.output.lower()
assert "00-of-01" in result.output
deleter.delete_batch.assert_not_called()

94
api/tests/unit_tests/libs/test_archive_storage.py
View File

@ -4,7 +4,7 @@ from datetime import datetime
from unittest.mock import ANY, MagicMock
import pytest
from botocore.exceptions import ClientError
from botocore.exceptions import ClientError, EndpointConnectionError
from libs import archive_storage as storage_module
from libs.archive_storage import (
@ -34,6 +34,10 @@ def _client_error(code: str) -> ClientError:
return ClientError({"Error": {"Code": code}}, "Operation")
def _network_error() -> EndpointConnectionError:
return EndpointConnectionError(endpoint_url="https://storage.example.com")
def _mock_client(monkeypatch: pytest.MonkeyPatch):
client = MagicMock()
client.head_bucket.return_value = None
@ -153,16 +157,38 @@ def test_get_object_returns_bytes(monkeypatch: pytest.MonkeyPatch):
assert storage.get_object("key") == b"payload"
def test_get_object_missing(monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"])
def test_get_object_missing(monkeypatch: pytest.MonkeyPatch, error_code: str):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.get_object.side_effect = _client_error("NoSuchKey")
client.get_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(FileNotFoundError, match="Archive object not found"):
storage.get_object("missing")
@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"])
def test_get_object_non_missing_error_fails_closed(monkeypatch: pytest.MonkeyPatch, error_code: str):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.get_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to download object"):
storage.get_object("key")
def test_get_object_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.get_object.side_effect = _network_error()
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to download object"):
storage.get_object("key")
def test_get_object_stream(monkeypatch: pytest.MonkeyPatch):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
@ -174,30 +200,80 @@ def test_get_object_stream(monkeypatch: pytest.MonkeyPatch):
assert list(storage.get_object_stream("key")) == [b"a", b"b"]
def test_get_object_stream_missing(monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"])
def test_get_object_stream_missing(monkeypatch: pytest.MonkeyPatch, error_code: str):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.get_object.side_effect = _client_error("NoSuchKey")
client.get_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(FileNotFoundError, match="Archive object not found"):
list(storage.get_object_stream("missing"))
def test_object_exists(monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"])
def test_object_exists_returns_false_only_for_not_found(
monkeypatch: pytest.MonkeyPatch,
error_code: str,
):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
storage = ArchiveStorage(bucket=BUCKET_NAME)
assert storage.object_exists("key") is True
client.head_object.side_effect = _client_error("404")
client.head_object.side_effect = _client_error(error_code)
assert storage.object_exists("missing") is False
def test_delete_object_error(monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"])
def test_object_exists_raises_when_existence_is_unknown(
monkeypatch: pytest.MonkeyPatch,
error_code: str,
):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.delete_object.side_effect = _client_error("500")
client.head_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to check archive object"):
storage.object_exists("key")
def test_object_exists_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.head_object.side_effect = _network_error()
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to check archive object"):
storage.object_exists("key")
@pytest.mark.parametrize("error_code", ["403", "429", "500", "SlowDown"])
def test_delete_object_error(monkeypatch: pytest.MonkeyPatch, error_code: str):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.delete_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to delete object"):
storage.delete_object("key")
@pytest.mark.parametrize("error_code", ["404", "NoSuchKey", "NotFound"])
def test_delete_object_missing_is_idempotent(monkeypatch: pytest.MonkeyPatch, error_code: str):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.delete_object.side_effect = _client_error(error_code)
storage = ArchiveStorage(bucket=BUCKET_NAME)
storage.delete_object("missing")
def test_delete_object_network_error_fails_closed(monkeypatch: pytest.MonkeyPatch):
_configure_storage(monkeypatch)
client, _ = _mock_client(monkeypatch)
client.delete_object.side_effect = _network_error()
storage = ArchiveStorage(bucket=BUCKET_NAME)
with pytest.raises(ArchiveStorageError, match="Failed to delete object"):

766
api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py Normal file
View File

@ -0,0 +1,766 @@
import json
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import MagicMock, call, patch
import pytest
from services.retention.workflow_run.bundle_archive_maintenance import (
ARCHIVED_TABLES,
ArchiveBundleCatalogEntry,
BundleManifest,
BundleOperationResult,
BundleReference,
WorkflowRunBundleArchiveMaintenance,
)
from services.retention.workflow_run.constants import (
ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME,
ARCHIVE_BUNDLE_DELETED_MARKER_NAME,
ARCHIVE_BUNDLE_FORMAT,
ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME,
ARCHIVE_BUNDLE_RESTORED_MARKER_NAME,
ARCHIVE_BUNDLE_SCHEMA_VERSION,
)
TENANT_ID = "1251fe32-c0c7-4fe2-a7bd-a8105267faf5"
CATALOG_ID = "019f63b7-5ca4-7681-9ce0-800283608f39"
BUNDLE_ID = "bundle-a"
def _table_records(
**overrides: list[dict[str, Any]],
) -> dict[str, list[dict[str, Any]]]:
records = {table_name: [] for table_name in ARCHIVED_TABLES}
records.update(overrides)
return records
def _catalog_entry(*, catalog_id: str = CATALOG_ID, shard: str = "00-of-01") -> ArchiveBundleCatalogEntry:
return ArchiveBundleCatalogEntry(
catalog_id=catalog_id,
tenant_id=TENANT_ID,
year=2025,
month=3,
shard=shard,
bundle_id=BUNDLE_ID,
workflow_run_count=0,
row_count=0,
archive_bytes=0,
)
def _manifest(
entry: ArchiveBundleCatalogEntry,
*,
bundle_id: str = BUNDLE_ID,
table_records: dict[str, list[dict[str, Any]]] | None = None,
) -> bytes:
object_prefix = WorkflowRunBundleArchiveMaintenance._catalog_object_prefix(entry)
records = table_records or _table_records()
tables = {
table_name: {
"row_count": len(records[table_name]),
"checksum": "",
"size_bytes": 0,
"object_key": f"{object_prefix}/{table_name}.parquet",
}
for table_name in (
"workflow_runs",
"workflow_app_logs",
"workflow_node_executions",
"workflow_node_execution_offload",
"workflow_pauses",
"workflow_pause_reasons",
"workflow_trigger_logs",
)
}
return json.dumps(
{
"schema_version": ARCHIVE_BUNDLE_SCHEMA_VERSION,
"archive_format": ARCHIVE_BUNDLE_FORMAT,
"tenant_id": entry.tenant_id,
"tenant_prefix": entry.tenant_id[0],
"year": entry.year,
"month": entry.month,
"shard": entry.shard,
"bundle_id": bundle_id,
"object_prefix": object_prefix,
"workflow_run_count": len(records["workflow_runs"]),
"workflow_node_execution_count": len(records["workflow_node_executions"]),
"tables": tables,
"run_ids": [str(record["id"]) for record in records["workflow_runs"]],
}
).encode()
def _session_factory(session: MagicMock) -> MagicMock:
factory = MagicMock()
factory.return_value.__enter__.return_value = session
return factory
def _bundle_reference(
entry: ArchiveBundleCatalogEntry,
*,
table_records: dict[str, list[dict[str, Any]]] | None = None,
) -> BundleReference:
manifest = cast(BundleManifest, json.loads(_manifest(entry, table_records=table_records)))
return BundleReference(
catalog=entry,
object_prefix=manifest["object_prefix"],
manifest_key=f"{manifest['object_prefix']}/manifest.json",
manifest_size_bytes=0,
manifest=manifest,
)
def _sample_archive_records() -> dict[str, list[dict[str, Any]]]:
return _table_records(
workflow_runs=[
{"id": "run-1", "status": "succeeded"},
{"id": "run-2", "status": "failed"},
],
workflow_app_logs=[
{"id": "app-log-1", "workflow_run_id": "run-1"},
],
workflow_node_executions=[
{"id": "node-1", "workflow_run_id": "run-1"},
],
workflow_node_execution_offload=[
{"id": "offload-1", "node_execution_id": "node-1"},
],
workflow_pauses=[
{"id": "pause-1", "workflow_run_id": "run-1"},
],
workflow_pause_reasons=[
{"id": "reason-1", "pause_id": "pause-1"},
],
workflow_trigger_logs=[
{"id": "trigger-1", "workflow_run_id": "run-1"},
],
)
def test_catalog_discovery_is_ordered_and_limited_before_storage_io() -> None:
entry = _catalog_entry()
bundle = SimpleNamespace(
id=entry.catalog_id,
tenant_id=entry.tenant_id,
year=entry.year,
month=entry.month,
shard=entry.shard,
bundle_id=entry.bundle_id,
workflow_run_count=entry.workflow_run_count,
row_count=entry.row_count,
archive_bytes=entry.archive_bytes,
)
session = MagicMock()
session.get.return_value = bundle
session.scalars.return_value = [bundle]
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, storage),
session_factory=cast(MagicMock, _session_factory(session)),
)
entries = maintenance._list_catalog_entries(
tenant_ids=[TENANT_ID],
target_year=2025,
target_month=3,
after_catalog_id=CATALOG_ID,
limit=2,
)
statement = session.scalars.call_args.args[0]
rendered = str(statement)
assert "workflow_run_archive_bundles.year" in rendered
assert "workflow_run_archive_bundles.month" in rendered
assert "workflow_run_archive_bundles.id >" in rendered
assert "ORDER BY workflow_run_archive_bundles.id ASC" in rendered
assert "LIMIT" in rendered
assert entries == [entry]
storage.list_objects.assert_not_called()
def test_catalog_discovery_filters_and_validates_the_requested_shard() -> None:
entry = _catalog_entry(shard="03-of-16")
bundle = SimpleNamespace(
id=entry.catalog_id,
tenant_id=entry.tenant_id,
year=entry.year,
month=entry.month,
shard=entry.shard,
bundle_id=entry.bundle_id,
workflow_run_count=entry.workflow_run_count,
row_count=entry.row_count,
archive_bytes=entry.archive_bytes,
)
session = MagicMock()
session.get.return_value = bundle
session.scalars.return_value = [bundle]
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
entries = maintenance._list_catalog_entries(
tenant_ids=None,
target_year=2025,
target_month=3,
after_catalog_id=CATALOG_ID,
limit=2,
shard="03-of-16",
)
statement = session.scalars.call_args.args[0]
rendered = str(statement)
assert "workflow_run_archive_bundles.shard =" in rendered
assert entries == [entry]
session.get.return_value = SimpleNamespace(
year=2025,
month=3,
tenant_id=TENANT_ID,
shard="04-of-16",
)
with pytest.raises(ValueError, match="requested archive shard"):
maintenance._list_catalog_entries(
tenant_ids=None,
target_year=2025,
target_month=3,
after_catalog_id=CATALOG_ID,
limit=2,
shard="03-of-16",
)
def test_catalog_shard_preflight_rejects_mixed_layout_before_delete() -> None:
session = MagicMock()
session.scalars.return_value = ["00-of-01"]
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
with pytest.raises(ValueError, match=r"unexpected shards.*00-of-01"):
maintenance.validate_catalog_shards(
target_year=2025,
target_month=3,
shard_total=16,
)
statement = session.scalars.call_args.args[0]
rendered = str(statement)
assert "workflow_run_archive_bundles.year" in rendered
assert "workflow_run_archive_bundles.month" in rendered
assert "workflow_run_archive_bundles.shard NOT IN" in rendered
def test_catalog_shard_preflight_accepts_an_expected_subset() -> None:
session = MagicMock()
session.scalars.return_value = []
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
maintenance.validate_catalog_shards(
target_year=2025,
target_month=3,
shard_total=16,
)
def test_catalog_shard_preflight_uses_requested_tenant_scope() -> None:
session = MagicMock()
session.scalars.return_value = []
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
maintenance.validate_catalog_shards(
target_year=2025,
target_month=3,
shard_total=16,
tenant_ids=[TENANT_ID],
)
statement = session.scalars.call_args.args[0]
assert "workflow_run_archive_bundles.tenant_id IN" in str(statement)
@pytest.mark.parametrize(
("cursor_bundle", "tenant_ids", "error_message"),
[
(None, None, "does not exist"),
(
SimpleNamespace(year=2024, month=3, tenant_id=TENANT_ID),
None,
"requested archive month",
),
(
SimpleNamespace(year=2025, month=3, tenant_id="other-tenant"),
[TENANT_ID],
"requested tenant scope",
),
],
)
def test_catalog_discovery_rejects_cursor_outside_requested_scope(
cursor_bundle: SimpleNamespace | None,
tenant_ids: list[str] | None,
error_message: str,
) -> None:
session = MagicMock()
session.get.return_value = cursor_bundle
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
with pytest.raises(ValueError, match=error_message):
maintenance._list_catalog_entries(
tenant_ids=tenant_ids,
target_year=2025,
target_month=3,
after_catalog_id=CATALOG_ID,
limit=1,
)
session.scalars.assert_not_called()
def test_catalog_manifest_identity_mismatch_fails_closed() -> None:
entry = _catalog_entry()
storage = MagicMock()
storage.get_object.return_value = _manifest(entry, bundle_id="other-bundle")
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, storage),
session_factory=cast(MagicMock, _session_factory(MagicMock())),
)
with pytest.raises(ValueError, match="identity does not match catalog"):
maintenance._build_bundle_reference(cast(MagicMock, storage), entry)
def test_bundle_maintenance_locks_the_existing_catalog_row() -> None:
entry = _catalog_entry()
session = MagicMock()
session.scalar.return_value = entry.catalog_id
WorkflowRunBundleArchiveMaintenance._lock_catalog_entry(session, entry)
statement = session.scalar.call_args.args[0]
rendered = str(statement)
assert "workflow_run_archive_bundles.id" in rendered
assert "workflow_run_archive_bundles.tenant_id" in rendered
assert "FOR UPDATE" in rendered
def test_failure_and_dry_run_do_not_return_a_persistable_cursor() -> None:
entry = _catalog_entry()
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, storage),
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_list_catalog_entries", return_value=[entry]),
patch.object(maintenance, "_build_bundle_reference", side_effect=RuntimeError("manifest unavailable")),
):
failed_summary = maintenance.delete_batch(
tenant_ids=None,
target_year=2025,
target_month=3,
after_catalog_id=None,
limit=1,
)
assert failed_summary.bundles_failed == 1
assert failed_summary.next_catalog_id is None
assert failed_summary.preview_next_catalog_id is None
dry_run = WorkflowRunBundleArchiveMaintenance(
dry_run=True,
storage=cast(MagicMock, storage),
session_factory=cast(MagicMock, _session_factory(session)),
)
bundle_ref = BundleReference(
catalog=entry,
object_prefix="object-prefix",
manifest_key="manifest.json",
manifest_size_bytes=0,
manifest=cast(BundleManifest, {}),
)
successful_result = BundleOperationResult(
catalog_id=entry.catalog_id,
bundle_id=entry.bundle_id,
tenant_id=entry.tenant_id,
object_prefix=bundle_ref.object_prefix,
success=True,
)
with (
patch.object(dry_run, "_list_catalog_entries", return_value=[entry]),
patch.object(dry_run, "_build_bundle_reference", return_value=bundle_ref),
patch.object(dry_run, "_delete_bundle", return_value=successful_result),
):
dry_run_summary = dry_run.delete_batch(
tenant_ids=None,
target_year=2025,
target_month=3,
after_catalog_id=None,
limit=1,
)
assert dry_run_summary.next_catalog_id is None
assert dry_run_summary.preview_next_catalog_id == entry.catalog_id
def test_live_archive_subset_accepts_full_partial_and_absent_live_data() -> None:
archive_records = _sample_archive_records()
manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest
partial_records = _table_records(
workflow_node_execution_offload=archive_records["workflow_node_execution_offload"],
workflow_pause_reasons=archive_records["workflow_pause_reasons"],
)
for live_records in (archive_records, partial_records, _table_records()):
WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset(
manifest,
archive_records,
live_records,
)
def test_live_archive_subset_rejects_extra_rows() -> None:
archive_records = _sample_archive_records()
manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest
live_records = _table_records(
workflow_app_logs=[
{"id": "extra-app-log", "workflow_run_id": "run-1"},
]
)
with pytest.raises(ValueError, match="rows missing from archive for workflow_app_logs"):
WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset(
manifest,
archive_records,
live_records,
)
def test_live_archive_subset_rejects_content_mismatch() -> None:
archive_records = _sample_archive_records()
manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest
live_records = _table_records(
workflow_runs=[
{"id": "run-1", "status": "failed"},
]
)
with pytest.raises(ValueError, match="subset content checksum mismatch for workflow_runs"):
WorkflowRunBundleArchiveMaintenance._validate_live_archive_subset(
manifest,
archive_records,
live_records,
)
def test_live_bundle_scope_includes_archived_ids_and_indirect_children() -> None:
archive_records = _sample_archive_records()
manifest = _bundle_reference(_catalog_entry(), table_records=archive_records).manifest
session = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
def select_live_parent_ids(_session, model, _run_ids):
if model.__tablename__ == "workflow_node_executions":
return ["live-node"]
if model.__tablename__ == "workflow_pauses":
return ["live-pause"]
raise AssertionError(f"unexpected model: {model}")
with (
patch.object(maintenance, "_select_ids_by_run_ids", side_effect=select_live_parent_ids),
patch.object(maintenance, "_load_records_by_column", return_value=[]) as load_records,
):
maintenance._load_live_bundle_records(
session,
manifest,
archive_records,
lock=True,
)
queries = [
(
call.args[1].__tablename__,
call.args[2].key,
set(call.args[3]),
call.kwargs["lock"],
)
for call in load_records.call_args_list
]
assert ("workflow_pause_reasons", "pause_id", {"pause-1", "live-pause"}, True) in queries
assert ("workflow_pause_reasons", "id", {"reason-1"}, True) in queries
assert ("workflow_node_execution_offload", "node_execution_id", {"node-1", "live-node"}, True) in queries
assert ("workflow_node_execution_offload", "id", {"offload-1"}, True) in queries
assert ("workflow_app_logs", "workflow_run_id", {"run-1", "run-2"}, True) in queries
assert ("workflow_app_logs", "id", {"app-log-1"}, True) in queries
def test_delete_bundle_accepts_matching_partial_rows_and_deletes_only_that_subset() -> None:
entry = _catalog_entry()
archive_records = _sample_archive_records()
bundle_ref = _bundle_reference(entry, table_records=archive_records)
partial_records = _table_records(
workflow_node_execution_offload=archive_records["workflow_node_execution_offload"],
workflow_pause_reasons=archive_records["workflow_pause_reasons"],
)
expected_deleted_counts = {table_name: len(partial_records[table_name]) for table_name in ARCHIVED_TABLES}
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_is_restore_started", return_value=False),
patch.object(maintenance, "_is_deleted", return_value=False),
patch.object(
maintenance,
"_validate_archive_object",
return_value=(bundle_ref.manifest, archive_records, 123),
),
patch.object(
maintenance,
"_load_live_bundle_records",
side_effect=[partial_records, _table_records()],
),
patch.object(
maintenance,
"_delete_bundle_rows",
return_value=expected_deleted_counts,
) as delete_bundle_rows,
patch.object(maintenance, "_put_marker"),
patch.object(maintenance, "_mark_deleted") as mark_deleted,
patch.object(maintenance, "_delete_marker"),
):
result = maintenance._delete_bundle(session, storage, bundle_ref)
assert result.success
delete_bundle_rows.assert_called_once_with(session, partial_records)
session.commit.assert_called_once_with()
session.rollback.assert_not_called()
mark_deleted.assert_called_once_with(storage, bundle_ref.object_prefix)
def test_delete_bundle_marks_an_already_absent_source_without_deleting_rows() -> None:
entry = _catalog_entry()
archive_records = _sample_archive_records()
bundle_ref = _bundle_reference(entry, table_records=archive_records)
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_is_restore_started", return_value=False),
patch.object(maintenance, "_is_deleted", return_value=False),
patch.object(
maintenance,
"_validate_archive_object",
return_value=(bundle_ref.manifest, archive_records, 123),
),
patch.object(maintenance, "_load_live_bundle_records", return_value=_table_records()),
patch.object(maintenance, "_delete_bundle_rows") as delete_bundle_rows,
patch.object(maintenance, "_mark_deleted") as mark_deleted,
patch.object(maintenance, "_delete_marker") as delete_marker,
):
result = maintenance._delete_bundle(session, storage, bundle_ref)
assert result.success
delete_bundle_rows.assert_not_called()
session.commit.assert_not_called()
session.rollback.assert_not_called()
mark_deleted.assert_called_once_with(storage, bundle_ref.object_prefix)
assert delete_marker.call_count == 2
def test_delete_bundle_with_deleted_marker_rejects_remaining_orphan_children() -> None:
entry = _catalog_entry()
archive_records = _sample_archive_records()
bundle_ref = _bundle_reference(entry, table_records=archive_records)
live_records = _table_records(
workflow_node_execution_offload=archive_records["workflow_node_execution_offload"],
)
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_is_restore_started", return_value=False),
patch.object(maintenance, "_is_deleted", return_value=True),
patch.object(
maintenance,
"_validate_archive_object",
return_value=(bundle_ref.manifest, archive_records, 123),
),
patch.object(maintenance, "_load_live_bundle_records", return_value=live_records),
patch.object(maintenance, "_delete_bundle_rows") as delete_bundle_rows,
):
result = maintenance._delete_bundle(session, storage, bundle_ref)
assert not result.success
assert "Live rows exist for bundle with deleted marker" in result.error
delete_bundle_rows.assert_not_called()
session.commit.assert_not_called()
session.rollback.assert_called_once_with()
def test_delete_bundle_rejects_an_in_progress_restore() -> None:
entry = _catalog_entry()
bundle_ref = _bundle_reference(entry)
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_is_restore_started", return_value=True),
patch.object(maintenance, "_validate_archive_object") as validate_archive,
):
result = maintenance._delete_bundle(session, storage, bundle_ref)
assert not result.success
assert "reconcile restore before delete" in result.error
validate_archive.assert_not_called()
session.commit.assert_not_called()
session.rollback.assert_called_once_with()
def test_delete_bundle_rows_use_only_verified_primary_keys() -> None:
live_records = _sample_archive_records()
session = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
session_factory=cast(MagicMock, _session_factory(session)),
)
with patch.object(
maintenance,
"_delete_by_column",
side_effect=lambda _session, _model, _column, values: len(values),
) as delete_by_column:
deleted_counts = maintenance._delete_bundle_rows(session, live_records)
expected_table_order = [
"workflow_pause_reasons",
"workflow_node_execution_offload",
"workflow_trigger_logs",
"workflow_app_logs",
"workflow_node_executions",
"workflow_pauses",
"workflow_runs",
]
assert [call.args[1].__tablename__ for call in delete_by_column.call_args_list] == expected_table_order
assert all(call.args[2].key == "id" for call in delete_by_column.call_args_list)
assert deleted_counts == {table_name: len(live_records[table_name]) for table_name in ARCHIVED_TABLES}
def test_restore_does_not_skip_an_interrupted_delete_without_deleted_marker() -> None:
entry = _catalog_entry()
session = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
with (
patch.object(maintenance, "_is_deleted", return_value=False),
patch.object(maintenance, "_is_delete_started", return_value=True),
patch.object(maintenance, "_validate_live_counts") as validate_live_counts,
):
result = maintenance._restore_bundle(session, MagicMock(), _bundle_reference(entry))
assert not result.success
assert "reconcile delete first" in result.error
validate_live_counts.assert_not_called()
session.commit.assert_not_called()
def test_restore_does_not_skip_missing_source_rows_without_deleted_marker() -> None:
entry = _catalog_entry()
session = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, MagicMock()),
session_factory=cast(MagicMock, _session_factory(session)),
)
bundle_ref = _bundle_reference(entry)
with (
patch.object(maintenance, "_is_deleted", return_value=False),
patch.object(maintenance, "_is_delete_started", return_value=False),
patch.object(
maintenance,
"_validate_live_counts",
side_effect=ValueError("source rows are missing"),
) as validate_live_counts,
):
result = maintenance._restore_bundle(session, MagicMock(), bundle_ref)
assert not result.success
assert "source rows are missing" in result.error
validate_live_counts.assert_called_once_with(session, bundle_ref.manifest, expected_present=True)
session.commit.assert_not_called()
def test_restore_reconciles_a_started_marker_after_the_source_commit() -> None:
entry = _catalog_entry()
session = MagicMock()
storage = MagicMock()
maintenance = WorkflowRunBundleArchiveMaintenance(
storage=cast(MagicMock, storage),
session_factory=cast(MagicMock, _session_factory(session)),
)
bundle_ref = _bundle_reference(entry)
with (
patch.object(maintenance, "_is_deleted", return_value=False),
patch.object(maintenance, "_is_delete_started", return_value=False),
patch.object(maintenance, "_is_restore_started", return_value=True),
patch.object(maintenance, "_validate_live_counts") as validate_live_counts,
patch.object(maintenance, "_mark_restored") as mark_restored,
):
result = maintenance._restore_bundle(session, storage, bundle_ref)
assert result.success
validate_live_counts.assert_called_once_with(session, bundle_ref.manifest, expected_present=True)
mark_restored.assert_called_once_with(storage, bundle_ref.object_prefix)
session.commit.assert_not_called()
def test_mark_restored_clears_stale_delete_marker_before_releasing_restore_fence() -> None:
storage = MagicMock()
object_prefix = "bundle-prefix"
operations = MagicMock()
with (
patch.object(WorkflowRunBundleArchiveMaintenance, "_delete_marker") as delete_marker,
patch.object(WorkflowRunBundleArchiveMaintenance, "_put_marker") as put_marker,
):
operations.attach_mock(delete_marker, "delete")
operations.attach_mock(put_marker, "put")
WorkflowRunBundleArchiveMaintenance._mark_restored(storage, object_prefix)
assert operations.mock_calls == [
call.delete(storage, object_prefix, ARCHIVE_BUNDLE_DELETED_MARKER_NAME),
call.put(storage, object_prefix, ARCHIVE_BUNDLE_RESTORED_MARKER_NAME),
call.delete(storage, object_prefix, ARCHIVE_BUNDLE_DELETE_STARTED_MARKER_NAME),
call.delete(storage, object_prefix, ARCHIVE_BUNDLE_RESTORE_STARTED_MARKER_NAME),
]