From 99e3b1a40147ed54a58694d76690efe239f788bd Mon Sep 17 00:00:00 2001 From: Ingram Z Date: Thu, 9 Jul 2026 17:46:28 +0800 Subject: [PATCH] fix: harden workflow archive DB retries (#38170) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/commands/retention.py | 158 ++++++++++++++---- .../archive_paid_plan_workflow_run.py | 101 +++++++++-- .../retention/workflow_run/db_retry.py | 66 ++++++++ .../retention/test_workflow_run_archiver.py | 118 +++++++++++++ .../commands/test_archive_workflow_runs.py | 139 +++++++++++++++ 5 files changed, 528 insertions(+), 54 deletions(-) create mode 100644 api/services/retention/workflow_run/db_retry.py create mode 100644 api/tests/unit_tests/commands/test_archive_workflow_runs.py diff --git a/api/commands/retention.py b/api/commands/retention.py index 3fe640c2f0c..841f18ac874 100644 --- a/api/commands/retention.py +++ b/api/commands/retention.py @@ -1,10 +1,12 @@ import datetime import logging import time +from collections.abc import Callable from typing import TypedDict import click import sqlalchemy as sa +from sqlalchemy.orm import Session, sessionmaker from extensions.ext_database import db from libs.datetime_utils import naive_utc_now @@ -12,6 +14,7 @@ from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpi from services.retention.conversation.messages_clean_policy import create_message_clean_policy from services.retention.conversation.messages_clean_service import MessagesCleanService from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup +from services.retention.workflow_run.db_retry import run_with_db_retry from services.retention.workflow_run.tenant_prefix import tenant_prefix_condition from tasks.remove_app_and_related_data_task import delete_draft_variables_batch @@ -35,6 +38,12 @@ class WorkflowRunArchiveTenantPlan(TypedDict): unpaid_tenant_ids: list[str] +class WorkflowRunArchivePrefixStats(TypedDict): + tenant_ids: list[str] + workflow_runs: int + workflow_node_executions: int + + def _normalize_utc_datetime(value: datetime.datetime) -> datetime.datetime: if value.tzinfo is None: return value.replace(tzinfo=datetime.UTC) @@ -57,6 +66,7 @@ def _parse_tenant_prefixes(prefixes: str | None) -> list[str]: def _get_archive_candidate_tenant_ids_by_prefix( + session: Session, prefix: str, *, start_from: datetime.datetime | None, @@ -75,7 +85,7 @@ def _get_archive_candidate_tenant_ids_by_prefix( if start_from is not None: conditions.append(WorkflowRun.created_at >= start_from) - tenant_ids = db.session.scalars( + tenant_ids = session.scalars( sa.select(WorkflowRun.tenant_id).where(*conditions).distinct().order_by(WorkflowRun.tenant_id) ).all() return list(tenant_ids) @@ -102,8 +112,80 @@ def _filter_paid_workflow_archive_tenant_ids(tenant_ids: list[str]) -> tuple[lis return paid_tenant_ids, unpaid_tenant_ids +def _run_archive_command_db_retry[T](operation_name: str, operation: Callable[[], T]) -> T: + return run_with_db_retry(operation_name, operation, logger=logger) + + +def _get_archive_candidate_tenant_ids_with_retry( + session_maker: sessionmaker[Session], + prefix: str, + *, + start_from: datetime.datetime | None, + end_before: datetime.datetime, +) -> list[str]: + def fetch_tenant_ids() -> list[str]: + with session_maker() as session: + return _get_archive_candidate_tenant_ids_by_prefix( + session, + prefix, + start_from=start_from, + end_before=end_before, + ) + + return _run_archive_command_db_retry(f"workflow archive tenant resolve for prefix {prefix}", fetch_tenant_ids) + + +def _get_archive_plan_prefix_stats( + session_maker: sessionmaker[Session], + prefix: str, + *, + start_from: datetime.datetime | None, + end_before: datetime.datetime, +) -> WorkflowRunArchivePrefixStats: + from graphon.enums import WorkflowExecutionStatus + from models.workflow import WorkflowNodeExecutionModel, WorkflowRun + from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver + + def fetch_prefix_stats() -> WorkflowRunArchivePrefixStats: + with session_maker() as session: + tenant_ids = _get_archive_candidate_tenant_ids_by_prefix( + session, + prefix, + start_from=start_from, + end_before=end_before, + ) + run_conditions = [ + WorkflowRun.created_at < end_before, + WorkflowRun.status.in_(WorkflowExecutionStatus.ended_values()), + WorkflowRun.type.in_(WorkflowRunArchiver.ARCHIVED_TYPE), + tenant_prefix_condition(WorkflowRun.tenant_id, prefix), + ] + if start_from is not None: + run_conditions.append(WorkflowRun.created_at >= start_from) + workflow_runs = ( + session.scalar(sa.select(sa.func.count()).select_from(WorkflowRun).where(*run_conditions)) or 0 + ) + candidate_runs = sa.select(WorkflowRun.id).where(*run_conditions).subquery() + workflow_node_executions = ( + session.scalar( + sa.select(sa.func.count()) + .select_from(WorkflowNodeExecutionModel) + .join(candidate_runs, WorkflowNodeExecutionModel.workflow_run_id == candidate_runs.c.id) + ) + or 0 + ) + return WorkflowRunArchivePrefixStats( + tenant_ids=tenant_ids, + workflow_runs=workflow_runs, + workflow_node_executions=workflow_node_executions, + ) + + return _run_archive_command_db_retry(f"workflow archive plan for prefix {prefix}", fetch_prefix_stats) + + def _resolve_archive_tenant_ids_from_plan( *, + session_maker: sessionmaker[Session], tenant_ids: str | None, tenant_prefixes: list[str], start_from: datetime.datetime | None, @@ -122,7 +204,8 @@ def _resolve_archive_tenant_ids_from_plan( requested_tenant_ids = [] for prefix in tenant_prefixes: requested_tenant_ids.extend( - _get_archive_candidate_tenant_ids_by_prefix( + _get_archive_candidate_tenant_ids_with_retry( + session_maker, prefix, start_from=start_from, end_before=end_before, @@ -143,6 +226,21 @@ def _resolve_archive_tenant_ids_from_plan( ) +def _safe_remove_scoped_session(context: str) -> None: + try: + db.session.remove() + except Exception: + logger.warning("Ignoring DB scoped-session cleanup error after %s", context, exc_info=True) + try: + db.session.registry.clear() + except Exception: + logger.warning("Ignoring DB scoped-session registry cleanup error after %s", context, exc_info=True) + try: + db.engine.dispose() + except Exception: + logger.warning("Ignoring DB engine dispose error after %s", context, exc_info=True) + + def _resolve_archive_time_range( *, before_days: int, @@ -349,10 +447,6 @@ def archive_workflow_runs_plan( supported workflow types, and the requested created_at window. V2 bundle archive does not maintain per-run archive logs, so this plan reports source-table volume. """ - from graphon.enums import WorkflowExecutionStatus - from models.workflow import WorkflowNodeExecutionModel, WorkflowRun - from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver - before_days, start_from, end_before = _resolve_archive_time_range( before_days=before_days, from_days_ago=from_days_ago, @@ -364,37 +458,25 @@ def archive_workflow_runs_plan( if include_archived: click.echo(click.style("--include-archived is a no-op for V2 bundle archive plans.", fg="yellow")) + session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) rows: list[WorkflowRunArchivePlanRow] = [] for prefix in _HEX_PREFIXES: - tenant_ids = _get_archive_candidate_tenant_ids_by_prefix( - prefix, - start_from=start_from, - end_before=plan_end_before, - ) + try: + prefix_stats = _get_archive_plan_prefix_stats( + session_maker, + prefix, + start_from=start_from, + end_before=plan_end_before, + ) + except Exception as exc: + logger.exception("Failed to build workflow archive plan for prefix %s", prefix) + raise click.ClickException(f"Failed to build workflow archive plan for prefix {prefix}.") from exc + tenant_ids = prefix_stats["tenant_ids"] + workflow_runs = prefix_stats["workflow_runs"] + workflow_node_executions = prefix_stats["workflow_node_executions"] total_tenants = len(tenant_ids) paid_tenant_ids, unpaid_tenant_ids = _filter_paid_workflow_archive_tenant_ids(tenant_ids) - run_conditions = [ - WorkflowRun.created_at < plan_end_before, - WorkflowRun.status.in_(WorkflowExecutionStatus.ended_values()), - WorkflowRun.type.in_(WorkflowRunArchiver.ARCHIVED_TYPE), - tenant_prefix_condition(WorkflowRun.tenant_id, prefix), - ] - if start_from is not None: - run_conditions.append(WorkflowRun.created_at >= start_from) - workflow_runs = ( - db.session.scalar(sa.select(sa.func.count()).select_from(WorkflowRun).where(*run_conditions)) or 0 - ) - candidate_runs = sa.select(WorkflowRun.id).where(*run_conditions).subquery() - workflow_node_executions = ( - db.session.scalar( - sa.select(sa.func.count()) - .select_from(WorkflowNodeExecutionModel) - .join(candidate_runs, WorkflowNodeExecutionModel.workflow_run_id == candidate_runs.c.id) - ) - or 0 - ) - rows.append( WorkflowRunArchivePlanRow( tenant_prefix=prefix, @@ -574,17 +656,18 @@ def archive_workflow_runs( ) ) + session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) try: tenant_plan = _resolve_archive_tenant_ids_from_plan( + session_maker=session_maker, tenant_ids=tenant_ids, tenant_prefixes=parsed_tenant_prefixes, start_from=start_from, end_before=plan_end_before, ) - except Exception: + except Exception as exc: logger.exception("Failed to resolve workflow archive tenant plan") - click.echo(click.style("Failed to resolve workflow archive tenant plan.", fg="red")) - return + raise click.ClickException("Failed to resolve workflow archive tenant plan.") from exc planned_tenant_ids = tenant_plan["archive_tenant_ids"] planned_paid_tenant_ids = tenant_plan["paid_tenant_ids"] if planned_tenant_ids is not None else None @@ -616,7 +699,10 @@ def archive_workflow_runs( dry_run=dry_run, delete_after_archive=delete_after_archive, ) - summary = archiver.run() + try: + summary = archiver.run() + finally: + _safe_remove_scoped_session("archive workflow run command") click.echo( click.style( f"Summary: processed={summary.total_runs_processed}, archived={summary.runs_archived}, " diff --git a/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py b/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py index db842a6ae48..af2bb023612 100644 --- a/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py +++ b/api/services/retention/workflow_run/archive_paid_plan_workflow_run.py @@ -29,12 +29,12 @@ import hashlib import json import logging import time -from collections.abc import Sequence +from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from enum import Enum from threading import Lock -from typing import Any, NotRequired, TypedDict, cast +from typing import Any, NotRequired, TypedDict, TypeVar, cast import click import pyarrow as pa @@ -70,8 +70,10 @@ from services.retention.workflow_run.constants import ( ARCHIVE_BUNDLE_MANIFEST_NAME, ARCHIVE_BUNDLE_SCHEMA_VERSION, ) +from services.retention.workflow_run.db_retry import is_retryable_db_disconnect, run_with_db_retry logger = logging.getLogger(__name__) +T = TypeVar("T") class TableStatsManifestEntry(TypedDict): @@ -208,6 +210,8 @@ class WorkflowRunArchiver: "workflow_pause_reasons", "workflow_trigger_logs", ] + DB_RETRY_ATTEMPTS = 3 + DB_RETRY_DELAYS_SECONDS = (1.0, 2.0) start_from: datetime.datetime | None end_before: datetime.datetime @@ -455,16 +459,40 @@ class WorkflowRunArchiver: """Fetch a batch of workflow runs to archive.""" repo = self._get_workflow_run_repo() tenant_ids = list(tenant_scope) if tenant_scope is not None else self.tenant_ids or None - return repo.get_runs_batch_by_time_range( - start_from=self.start_from, - end_before=self.end_before, - last_seen=last_seen, - batch_size=self.batch_size, - run_types=self.ARCHIVED_TYPE, - tenant_ids=tenant_ids, - tenant_prefixes=None if tenant_ids else self.tenant_prefixes or None, - run_shard_index=self.run_shard_index, - run_shard_total=self.run_shard_total, + + return self._run_with_db_retry( + "workflow run batch fetch", + lambda: repo.get_runs_batch_by_time_range( + start_from=self.start_from, + end_before=self.end_before, + last_seen=last_seen, + batch_size=self.batch_size, + run_types=self.ARCHIVED_TYPE, + tenant_ids=tenant_ids, + tenant_prefixes=None if tenant_ids else self.tenant_prefixes or None, + run_shard_index=self.run_shard_index, + run_shard_total=self.run_shard_total, + ), + ) + + @staticmethod + def _is_retryable_db_disconnect(exc: BaseException) -> bool: + return is_retryable_db_disconnect(exc) + + @staticmethod + def _safe_rollback(session: Session, bundle_id: str) -> None: + try: + session.rollback() + except Exception: + logger.warning("Failed to rollback archive session for bundle %s", bundle_id, exc_info=True) + + def _run_with_db_retry(self, operation_name: str, operation: Callable[[], T]) -> T: + return run_with_db_retry( + operation_name, + operation, + logger=logger, + attempts=self.DB_RETRY_ATTEMPTS, + delays_seconds=self.DB_RETRY_DELAYS_SECONDS, ) def _tenant_scan_scopes(self) -> list[list[str] | None]: @@ -532,16 +560,14 @@ class WorkflowRunArchiver: if self.workers == 1 or len(bundle_groups) == 1: results: list[ArchiveResult] = [] for bundle_runs in bundle_groups: - with session_maker() as session: - results.append(self._archive_bundle(session, storage, bundle_runs)) + results.append(self._archive_bundle_with_retry(session_maker, storage, bundle_runs)) return results results = [] max_workers = min(self.workers, len(bundle_groups)) def archive_in_worker(bundle_runs: Sequence[WorkflowRun]) -> ArchiveResult: - with session_maker() as session: - return self._archive_bundle(session, storage, bundle_runs) + return self._archive_bundle_with_retry(session_maker, storage, bundle_runs) with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [executor.submit(archive_in_worker, bundle_runs) for bundle_runs in bundle_groups] @@ -549,6 +575,39 @@ class WorkflowRunArchiver: results.append(future.result()) return results + def _archive_bundle_with_retry( + self, + session_maker: sessionmaker[Session], + storage: ArchiveStorage | None, + runs: Sequence[WorkflowRun], + ) -> ArchiveResult: + identity = self._build_bundle_identity(runs) + + try: + return self._run_with_db_retry( + f"archive workflow run bundle {identity.bundle_id}", + lambda: self._archive_bundle_once(session_maker, storage, runs), + ) + except Exception as exc: + logger.exception("Failed to archive workflow run bundle %s after retries", identity.bundle_id) + return ArchiveResult( + bundle_id=identity.bundle_id, + tenant_id=identity.tenant_id, + object_prefix=identity.object_prefix, + run_count=len(runs), + success=False, + error=str(exc), + ) + + def _archive_bundle_once( + self, + session_maker: sessionmaker[Session], + storage: ArchiveStorage | None, + runs: Sequence[WorkflowRun], + ) -> ArchiveResult: + with session_maker() as session: + return self._archive_bundle(session, storage, runs) + def _archive_bundle( self, session: Session, @@ -645,9 +704,12 @@ class WorkflowRunArchiver: result.success = True except Exception as e: + if self._is_retryable_db_disconnect(e): + self._safe_rollback(session, identity.bundle_id) + raise logger.exception("Failed to archive workflow run bundle %s", identity.bundle_id) result.error = str(e) - session.rollback() + self._safe_rollback(session, identity.bundle_id) result.elapsed_time = time.time() - start_time return result @@ -794,6 +856,9 @@ class WorkflowRunArchiver: end_before = self.end_before if end_before is None: raise ValueError("archive window end must be set") + formatted_end_before = self._format_window_datetime(end_before) + if formatted_end_before is None: + raise ValueError("archive window end must be set") return ArchiveManifestDict( schema_version=ARCHIVE_BUNDLE_SCHEMA_VERSION, archive_format=ARCHIVE_BUNDLE_FORMAT, @@ -813,7 +878,7 @@ class WorkflowRunArchiver: archived_at=datetime.datetime.now(datetime.UTC).isoformat(), campaign_id=self.campaign_id, archive_window_start=self._format_window_datetime(self.start_from), - archive_window_end=end_before.isoformat(), + archive_window_end=formatted_end_before, run_shard=identity.shard, tables=tables, run_ids=[run.id for run in sorted_runs], diff --git a/api/services/retention/workflow_run/db_retry.py b/api/services/retention/workflow_run/db_retry.py new file mode 100644 index 00000000000..d39adb37a0b --- /dev/null +++ b/api/services/retention/workflow_run/db_retry.py @@ -0,0 +1,66 @@ +import logging +import time +from collections.abc import Callable + +from sqlalchemy.exc import DBAPIError +from sqlalchemy.exc import OperationalError as SQLAlchemyOperationalError + +DEFAULT_DB_RETRY_ATTEMPTS = 3 +DEFAULT_DB_RETRY_DELAYS_SECONDS = (1.0, 2.0) + +_DB_DISCONNECT_PATTERNS = ( + "server closed the connection unexpectedly", + "connection already closed", + "closed the connection", + "connection not open", + "terminating connection", + "connection reset", + "broken pipe", + "connection invalidated", +) + + +def is_retryable_db_disconnect(exc: BaseException) -> bool: + if isinstance(exc, DBAPIError) and exc.connection_invalidated: + return True + + if not _is_db_operational_error(exc): + return False + + original_exception = exc.orig if isinstance(exc, DBAPIError) else None + message = f"{exc} {original_exception or ''}".lower() + return any(pattern in message for pattern in _DB_DISCONNECT_PATTERNS) + + +def run_with_db_retry[T]( + operation_name: str, + operation: Callable[[], T], + *, + logger: logging.Logger, + attempts: int = DEFAULT_DB_RETRY_ATTEMPTS, + delays_seconds: tuple[float, ...] = DEFAULT_DB_RETRY_DELAYS_SECONDS, +) -> T: + for attempt in range(1, attempts + 1): + try: + return operation() + except Exception as exc: + if not is_retryable_db_disconnect(exc) or attempt == attempts: + raise + delay = delays_seconds[min(attempt - 1, len(delays_seconds) - 1)] + logger.warning( + "Retrying %s after retryable DB disconnect (attempt %s/%s, sleep %.1fs)", + operation_name, + attempt, + attempts, + delay, + exc_info=True, + ) + time.sleep(delay) + raise RuntimeError(f"{operation_name} did not complete") + + +def _is_db_operational_error(exc: BaseException) -> bool: + if isinstance(exc, SQLAlchemyOperationalError): + return True + + return exc.__class__.__name__ == "OperationalError" and exc.__class__.__module__.startswith(("psycopg", "psycopg2")) diff --git a/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py b/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py index 5bea9469e12..4d6ff7bd7ec 100644 --- a/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py +++ b/api/tests/integration_tests/services/retention/test_workflow_run_archiver.py @@ -6,8 +6,10 @@ from unittest.mock import ANY, MagicMock, patch import pyarrow as pa import pyarrow.parquet as pq import pytest +from sqlalchemy.exc import OperationalError from services.retention.workflow_run.archive_paid_plan_workflow_run import ( + ArchiveResult, ArchiveSummary, WorkflowRunArchiver, ) @@ -32,6 +34,30 @@ class FakeArchiveStorage: return sorted(key for key in self.objects if key.startswith(prefix)) +def _db_disconnect_error() -> OperationalError: + return OperationalError( + "select 1", + {}, + RuntimeError("server closed the connection unexpectedly"), + connection_invalidated=True, + ) + + +def _run(run_id: str = "run-1"): + run = MagicMock() + run.id = run_id + run.tenant_id = "tenant-1" + run.created_at = datetime.datetime(2025, 3, 15, 10, 0, 0) + return run + + +def _session_context(session): + context = MagicMock() + context.__enter__.return_value = session + context.__exit__.return_value = False + return context + + class TestWorkflowRunArchiverInit: def test_start_from_without_end_before_raises(self): with pytest.raises(ValueError, match="start_from and end_before must be provided together"): @@ -139,6 +165,32 @@ class TestWorkflowRunArchiverInit: repo.get_runs_batch_by_time_range.assert_called_once() assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["tenant-b"] + def test_get_runs_batch_retries_retryable_db_disconnect(self): + repo = MagicMock() + repo.get_runs_batch_by_time_range.side_effect = [_db_disconnect_error(), []] + archiver = WorkflowRunArchiver(workflow_run_repo=repo) + + with patch("services.retention.workflow_run.db_retry.time.sleep") as sleep: + runs = archiver._get_runs_batch(None) + + assert runs == [] + assert repo.get_runs_batch_by_time_range.call_count == 2 + sleep.assert_called_once_with(1.0) + + def test_get_runs_batch_does_not_retry_non_db_broken_pipe_error(self): + repo = MagicMock() + repo.get_runs_batch_by_time_range.side_effect = RuntimeError("broken pipe") + archiver = WorkflowRunArchiver(workflow_run_repo=repo) + + with ( + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + pytest.raises(RuntimeError, match="broken pipe"), + ): + archiver._get_runs_batch(None) + + repo.get_runs_batch_by_time_range.assert_called_once() + sleep.assert_not_called() + def test_start_message_includes_shard(self): archiver = WorkflowRunArchiver(tenant_prefixes=["0"], run_shard_index=1, run_shard_total=4) @@ -351,6 +403,72 @@ class TestDryRunArchive: assert summary.table_stats["workflow_app_logs"].size_bytes == 32 +class TestArchiveDbRetry: + def test_archive_bundle_groups_retries_with_fresh_session(self): + archiver = WorkflowRunArchiver(days=90) + run = _run() + session_maker = MagicMock( + side_effect=[ + _session_context(MagicMock(name="session-1")), + _session_context(MagicMock(name="session-2")), + ] + ) + success = ArchiveResult( + bundle_id=archiver._build_bundle_identity([run]).bundle_id, + tenant_id=run.tenant_id, + object_prefix=archiver._build_bundle_identity([run]).object_prefix, + run_count=1, + success=True, + ) + + with ( + patch.object(archiver, "_archive_bundle", side_effect=[_db_disconnect_error(), success]) as archive_bundle, + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + ): + results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) + + assert results == [success] + assert archive_bundle.call_count == 2 + assert session_maker.call_count == 2 + sleep.assert_called_once_with(1.0) + + def test_archive_bundle_groups_returns_failed_result_after_retry_exhaustion(self): + archiver = WorkflowRunArchiver(days=90) + run = _run() + session_maker = MagicMock( + side_effect=[ + _session_context(MagicMock(name="session-1")), + _session_context(MagicMock(name="session-2")), + _session_context(MagicMock(name="session-3")), + ] + ) + + with ( + patch.object(archiver, "_archive_bundle", side_effect=[_db_disconnect_error()] * 3) as archive_bundle, + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + ): + results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) + + assert len(results) == 1 + assert results[0].success is False + assert "server closed the connection unexpectedly" in (results[0].error or "") + assert archive_bundle.call_count == archiver.DB_RETRY_ATTEMPTS + assert session_maker.call_count == archiver.DB_RETRY_ATTEMPTS + assert sleep.call_count == archiver.DB_RETRY_ATTEMPTS - 1 + + def test_archive_bundle_uses_safe_rollback_when_failure_rolls_back_badly(self): + archiver = WorkflowRunArchiver(days=90, dry_run=True) + session = MagicMock() + session.rollback.side_effect = RuntimeError("rollback failed") + + with patch.object(archiver, "_extract_bundle_data", side_effect=RuntimeError("extract failed")): + result = archiver._archive_bundle(session, None, [_run()]) + + assert result.success is False + assert result.error == "extract failed" + session.rollback.assert_called_once() + + class TestArchiveRunIdempotency: def _index_payload(self, archiver: WorkflowRunArchiver, run_ids: list[str], run) -> tuple[str, bytes]: identity = archiver._build_bundle_identity([run]) diff --git a/api/tests/unit_tests/commands/test_archive_workflow_runs.py b/api/tests/unit_tests/commands/test_archive_workflow_runs.py new file mode 100644 index 00000000000..2e2b5df8029 --- /dev/null +++ b/api/tests/unit_tests/commands/test_archive_workflow_runs.py @@ -0,0 +1,139 @@ +import datetime +from unittest.mock import MagicMock + +import click +import pytest +from sqlalchemy.exc import OperationalError + +from commands import retention + + +def _db_disconnect_error() -> OperationalError: + return OperationalError( + "select 1", + {}, + RuntimeError("server closed the connection unexpectedly"), + connection_invalidated=True, + ) + + +def _session_context(session): + context = MagicMock() + context.__enter__.return_value = session + context.__exit__.return_value = False + return context + + +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")] + session_maker = MagicMock(side_effect=[_session_context(sessions[0]), _session_context(sessions[1])]) + calls = [] + + def get_candidate_tenants(session, prefix, *, start_from, end_before): + calls.append((session, prefix, start_from, end_before)) + return [f"{prefix}-paid", f"{prefix}-free"] + + monkeypatch.setattr(retention, "_get_archive_candidate_tenant_ids_by_prefix", get_candidate_tenants) + monkeypatch.setattr( + retention, + "_filter_paid_workflow_archive_tenant_ids", + lambda tenant_ids: (["a-paid", "b-paid"], ["a-free", "b-free"]), + ) + + tenant_plan = retention._resolve_archive_tenant_ids_from_plan( + session_maker=session_maker, + tenant_ids=None, + tenant_prefixes=["a", "b"], + start_from=None, + end_before=end_before, + ) + + assert tenant_plan["archive_tenant_ids"] == ["a-paid", "b-paid"] + assert tenant_plan["paid_tenant_ids"] == ["a-paid", "b-paid"] + assert tenant_plan["unpaid_tenant_ids"] == ["a-free", "b-free"] + assert calls == [ + (sessions[0], "a", None, end_before), + (sessions[1], "b", None, end_before), + ] + + +def test_safe_remove_scoped_session_discards_registry_and_disposes_after_remove_error(monkeypatch): + fake_db = MagicMock() + fake_db.session.remove.side_effect = RuntimeError("server closed the connection unexpectedly") + monkeypatch.setattr(retention, "db", fake_db) + + retention._safe_remove_scoped_session("archive workflow run command") + + fake_db.session.remove.assert_called_once() + fake_db.session.registry.clear.assert_called_once() + fake_db.engine.dispose.assert_called_once() + + +def test_archive_command_db_retry_retries_retryable_db_disconnect(monkeypatch): + operation = MagicMock(side_effect=[_db_disconnect_error(), "ok"]) + sleep = MagicMock() + monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) + + result = retention._run_archive_command_db_retry("archive plan", operation) + + assert result == "ok" + assert operation.call_count == 2 + sleep.assert_called_once_with(1.0) + + +def test_archive_plan_prefix_stats_retries_count_query_with_fresh_session(monkeypatch): + end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC) + sessions = [MagicMock(name="session-1"), MagicMock(name="session-2")] + sessions[0].scalar.side_effect = _db_disconnect_error() + sessions[1].scalar.side_effect = [7, 9] + session_maker = MagicMock(side_effect=[_session_context(sessions[0]), _session_context(sessions[1])]) + sleep = MagicMock() + + monkeypatch.setattr( + retention, + "_get_archive_candidate_tenant_ids_by_prefix", + lambda session, prefix, *, start_from, end_before: [f"{prefix}-tenant"], + ) + monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) + + stats = retention._get_archive_plan_prefix_stats( + session_maker, + "a", + start_from=None, + end_before=end_before, + ) + + assert stats["tenant_ids"] == ["a-tenant"] + assert stats["workflow_runs"] == 7 + assert stats["workflow_node_executions"] == 9 + assert session_maker.call_count == 2 + sleep.assert_called_once_with(1.0) + + +def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(monkeypatch): + fake_db = MagicMock() + monkeypatch.setattr(retention, "db", fake_db) + monkeypatch.setattr( + retention, + "_resolve_archive_tenant_ids_from_plan", + MagicMock(side_effect=RuntimeError("tenant plan failed")), + ) + + with pytest.raises(click.ClickException, match="Failed to resolve workflow archive tenant plan"): + retention.archive_workflow_runs.callback( + tenant_ids="tenant-1", + tenant_prefixes=None, + before_days=90, + from_days_ago=None, + to_days_ago=None, + start_from=None, + end_before=None, + batch_size=10000, + workers=1, + run_shard_index=None, + run_shard_total=None, + limit=None, + dry_run=True, + delete_after_archive=False, + )