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

fix: harden workflow archive DB retries (#38170)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Ingram Z 2026-07-09 17:46:28 +08:00 committed by GitHub
parent 45957225cd
commit 99e3b1a401
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 528 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

158
api/commands/retention.py
View File

@ -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}, "

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

@ -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],

66
api/services/retention/workflow_run/db_retry.py Normal file
View File

@ -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"))

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

@ -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])

139
api/tests/unit_tests/commands/test_archive_workflow_runs.py Normal file
View File

@ -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,
)