From be26a57c1eec5b8b4b0bb32b344b5dd1ea4309c9 Mon Sep 17 00:00:00 2001 From: Danny <64348131+Danny5487401@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:46:50 +0800 Subject: [PATCH] feat: add bounded PostgreSQL workflow history retention --- .github/workflows/test-workflow-retention.yml | 61 +++ docs/deployment/configuration.mdx | 2 + docs/deployment/workflow-retention.mdx | 93 ++++ docs/mint.json | 1 + keep/api/core/workflow_retention.py | 511 +++++++++++++++++ .../versions/2026-11-12-13-30_6942a1b2c3d4.py | 86 +++ keep/api/models/db/workflow.py | 6 + keep/cli/workflow_retention.py | 98 ++++ scripts/cleanup_workflow_executions.py | 13 + tests/maintenance/conftest.py | 132 +++++ tests/maintenance/test_workflow_retention.py | 517 ++++++++++++++++++ .../test_workflow_retention_migration.py | 127 +++++ 12 files changed, 1647 insertions(+) create mode 100644 .github/workflows/test-workflow-retention.yml create mode 100644 docs/deployment/workflow-retention.mdx create mode 100644 keep/api/core/workflow_retention.py create mode 100644 keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py create mode 100644 keep/cli/workflow_retention.py create mode 100644 scripts/cleanup_workflow_executions.py create mode 100644 tests/maintenance/conftest.py create mode 100644 tests/maintenance/test_workflow_retention.py create mode 100644 tests/maintenance/test_workflow_retention_migration.py diff --git a/.github/workflows/test-workflow-retention.yml b/.github/workflows/test-workflow-retention.yml new file mode 100644 index 0000000000..5904b325c9 --- /dev/null +++ b/.github/workflows/test-workflow-retention.yml @@ -0,0 +1,61 @@ +name: PostgreSQL workflow retention tests +on: + pull_request: + paths: + - "keep/api/core/workflow_retention.py" + - "keep/cli/workflow_retention.py" + - "keep/api/models/db/workflow.py" + - "keep/api/models/db/migrations/**" + - "scripts/cleanup_workflow_executions.py" + - "tests/maintenance/**" + - ".github/workflows/test-workflow-retention.yml" + push: + branches: [main] + paths: + - "keep/api/core/workflow_retention.py" + - "keep/cli/workflow_retention.py" + - "keep/api/models/db/workflow.py" + - "keep/api/models/db/migrations/**" + - "scripts/cleanup_workflow_executions.py" + - "tests/maintenance/**" + - ".github/workflows/test-workflow-retention.yml" + workflow_dispatch: +permissions: + contents: read +jobs: + postgres: + runs-on: ubuntu-latest + strategy: + matrix: + python: ["3.11", "3.13"] + services: + postgres: + image: postgres:16.14-bookworm + env: + POSTGRES_DB: keep_retention_tests + POSTGRES_USER: keep_tests + POSTGRES_PASSWORD: keep_tests + ports: + - 5432:5432 + options: >- + --health-cmd="pg_isready -U keep_tests -d keep_retention_tests" + --health-interval=5s --health-timeout=5s --health-retries=10 + env: + KEEP_RETENTION_TEST_DATABASE_URL: postgresql+psycopg://keep_tests:keep_tests@localhost:5432/keep_retention_tests + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + - name: Install focused test dependencies + run: >- + python -m pip install + sqlalchemy==2.0.36 sqlmodel==0.0.22 pydantic==1.10.27 + alembic==1.13.3 "psycopg[binary]==3.2.3" pytest==8.3.5 + ruff==0.11.4 black==24.10.0 + - name: Lint and format changed Python + run: | + python -m ruff check keep/api/core/workflow_retention.py keep/cli/workflow_retention.py scripts/cleanup_workflow_executions.py tests/maintenance keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py + python -m black --check keep/api/core/workflow_retention.py keep/cli/workflow_retention.py scripts/cleanup_workflow_executions.py tests/maintenance keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py + - name: Test PostgreSQL constraints, bounded cleanup, and migration + run: python -m pytest --confcutdir=tests/maintenance tests/maintenance -q diff --git a/docs/deployment/configuration.mdx b/docs/deployment/configuration.mdx index 7738c7358b..3a7cbdff9c 100644 --- a/docs/deployment/configuration.mdx +++ b/docs/deployment/configuration.mdx @@ -64,6 +64,8 @@ Keep is highly configurable through environment variables. This allows you to cu | **DB_IP_TYPE** | Specifies the Cloud SQL IP type | No | "public" | "public", "private" or "psc" | | **SKIP_DB_CREATION** | Skips database creation and migrations | No | "false" | "true" or "false" | +For opt-in PostgreSQL execution/log history cleanup, see [Workflow history retention](/deployment/workflow-retention). + ### Resource Provisioning diff --git a/docs/deployment/workflow-retention.mdx b/docs/deployment/workflow-retention.mdx new file mode 100644 index 0000000000..ad12c3760c --- /dev/null +++ b/docs/deployment/workflow-retention.mdx @@ -0,0 +1,93 @@ +--- +title: "PostgreSQL workflow history retention" +sidebarTitle: "Workflow history retention" +--- + +Workflow execution results (including their TOAST storage) and execution logs can grow independently of alert history. The opt-in maintenance command below removes old terminal execution history; it does not delete alerts, incidents, workflow definitions, or versions. No background cleanup is enabled automatically. + +## Prepare the database + +This command currently supports PostgreSQL. Use an image/package containing `keep.cli.workflow_retention` and the same database configuration as the Keep API. The command reuses Keep's configured synchronous engine. Do not put database passwords in command-line arguments or shell history. + +Take a database backup and verify a restore procedure before the first apply. The first cleanup can generate substantial I/O and WAL; budget free space for index creation, temporary sort files, WAL, and ongoing application writes. + +The Alembic revision `6942a1b2c3d4` adds an index on `workflowexecutionlog.workflow_execution_id`. PostgreSQL builds it concurrently, outside the migration transaction. Existing large installations should prebuild it during a planned window **before upgrading** so API startup does not have to wait for a long index build: + +```sql +CREATE INDEX CONCURRENTLY idx_workflowexecutionlog_workflow_execution_id +ON workflowexecutionlog (workflow_execution_id); +``` + +Run this statement outside a transaction block. Check the existing definition before creating an index with that name. The migration accepts an existing valid, non-partial index beginning with `workflow_execution_id`; a failed build that left the named index invalid requires explicit operator inspection/repair before retrying. The migration never silently drops an invalid index. A downgrade removes only the named index, not equivalent indexes with other names. + +Cleanup performs no DDL and refuses to run without a valid, non-partial leading execution-ID index on the log and both association tables. It also refuses unknown inbound execution dependencies. These checks prevent a new schema dependency or an unindexed lookup from being silently ignored. + +## Preview first + +Find the tenant ID using the database's tenant metadata or your existing administration tools. You must explicitly choose `--tenant-id` or `--all-tenants`. + +```bash +python -m keep.cli.workflow_retention \ + --tenant-id TENANT_ID \ + --success-retention-days 14 \ + --failure-retention-days 30 \ + --success-log-retention-days 7 +``` + +From a source checkout, `python scripts/cleanup_workflow_executions.py` is an equivalent entry point. The packaged module is available in API images; it does not depend on a top-level `scripts` directory being included in the image. + +The default mode is read-only dry-run; `--dry-run` is an optional explicit spelling. Previews count the complete eligible scope instead of stopping after one batch, and include stored result bytes and log/association counts. Result bytes are not a prediction of filesystem space returned. A bounded, repeatable-read transaction provides a consistent preview. An incomplete/timeout report has `completed: false`; never interpret a partial count as the full deletion scope. + +## Apply bounded batches + +Use the same policy and scope, adding `--apply` only after checking the preview: + +```bash +python -m keep.cli.workflow_retention \ + --tenant-id TENANT_ID \ + --success-retention-days 14 \ + --failure-retention-days 30 \ + --success-log-retention-days 7 \ + --batch-size 100 \ + --log-batch-size 5000 \ + --max-runtime-seconds 600 \ + --apply +``` + +Defaults and limits: + +| Option | Default | Meaning | +|---|---:|---| +| `--success-retention-days` / `--retention-days` | 30 | Success execution/result retention | +| `--failure-retention-days` | 30 | Error and `providers_not_configured` retention | +| `--success-log-retention-days` | unset | Optional shorter success-log retention; must not exceed success retention | +| `--keep-failed` | off | Preserve failures indefinitely instead of applying failure retention | +| `--keep-successes` | 5 | Preserve the latest 1–100 non-test successes per tenant/workflow | +| `--batch-size` | 100 | At most 1,000 candidate parents per batch | +| `--log-batch-size` | 5,000 | At most 10,000 rows per child-table DELETE; association fan-out uses the same bound | +| `--batch-pause-ms` | 200 | Pause between committed batches; 0–10,000 milliseconds | +| `--max-runtime-seconds` | 600 | Runtime budget, checked between batches; maximum 3,600 | +| `--statement-timeout-ms` | 15,000 | Per-statement timeout, also limited by remaining budget at batch start | +| `--lock-timeout-ms` | 1,000 | Bound waits for database locks | + +Only `success`, `error`, and `providers_not_configured` executions with a valid nonnegative `execution_time` are eligible. Both the start and calculated completion time must be older than the relevant cutoff. Active, unknown, and ambiguous timeout executions are skipped. `is_running` is not treated as a boolean: Keep uses it as a scheduling sentinel/random integer. + +The latest non-test successful records and the scheduler-relevant last completed record per tenant/workflow are protected independently of age. Test executions are never used as checkpoints. Incremental workflows retain their complete protected results. The command does not rewrite checkpoints or change Keep's existing one-day successful-history lookup window; recovery after a longer pause still needs the documented workflow-specific recovery procedure. + +Children are drained in bounded transactions before deleting the execution. Large log or association sets can therefore be partially cleaned while their old parent remains for the next batch. Final parent deletion requires no remaining children, rechecks eligibility, and shares a transaction with that batch's association deletions. Earlier committed batches are preserved if a later batch fails; the current transaction rolls back. Apply counters describe committed rows only. Retrying with the same policy is supported. + +Exit codes: `0` means the scan completed, `2` means an incomplete report/runtime or query limit (and also argparse usage errors), and `1` means preflight or execution failure. Apply errors must be investigated before retrying. Counts and payloads are JSON on stdout; progress logs use stderr. No result payloads, SQL text, or connection credentials are included in the report. + +## Catch up and schedule + +Start with a small apply window and observe query latency, free space, WAL, application workflow success, and checkpoint continuity. Repeat bounded catch-up runs until the backlog is manageable, then schedule regular maintenance using your deployment's existing scheduler. Run one maintenance process per database; for Kubernetes CronJobs set `concurrencyPolicy: Forbid`, and do not overlap manual runs. The command is restartable but does not provide a cross-process lease or exactly-once scheduling guarantee. + +Retained history is no longer available in execution-history screens after deletion; alert and incident records remain. A code rollback cannot restore committed deletions. Restore historical records from a backup into a separate database and recover selectively instead of overwriting newer application data. + +## Vacuum and monitor + +After the initial catch-up, inspect dead tuples and autovacuum activity for both tables and their TOAST storage. Consider lower table-level autovacuum thresholds after measuring I/O: default scale factors can require millions of deleted rows on a large log table. A controlled ordinary `VACUUM (ANALYZE)` can help after a large cleanup, outside the cleanup transaction. + +DELETE and ordinary VACUUM generally make space reusable within PostgreSQL. They do not promise an immediate decrease in PVC filesystem usage. `VACUUM FULL` rewrites a table, needs extra space, and takes an exclusive lock; plan it separately if filesystem compaction is necessary. This command never performs VACUUM or truncation. + +Monitor the age and count of eligible history, committed cleanup throughput, incomplete/failed runs, free space and growth, WAL, and the protected workflows' checkpoint continuity. Reducing newly persisted intermediate outputs is a separate change; do not clear a successful checkpoint just to reduce history size. diff --git a/docs/mint.json b/docs/mint.json index e7fa63df66..ca60cf3960 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -278,6 +278,7 @@ "pages": [ "deployment/configuration", "deployment/monitoring", + "deployment/workflow-retention", { "group": "Authentication", "pages": [ diff --git a/keep/api/core/workflow_retention.py b/keep/api/core/workflow_retention.py new file mode 100644 index 0000000000..985043eafa --- /dev/null +++ b/keep/api/core/workflow_retention.py @@ -0,0 +1,511 @@ +"""Opt-in PostgreSQL workflow-history maintenance, independent of the scheduler.""" + +import logging +import time +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone + +from sqlalchemy import and_, delete, exists, func, or_, select, text +from sqlalchemy.exc import DBAPIError + +from keep.api.models.db.workflow import ( + WorkflowExecution, + WorkflowExecutionLog, + WorkflowToAlertExecution, + WorkflowToIncidentExecution, +) + +logger = logging.getLogger(__name__) +EXECUTIONS = WorkflowExecution.__table__ +LOGS = WorkflowExecutionLog.__table__ +ALERT_LINKS = WorkflowToAlertExecution.__table__ +INCIDENT_LINKS = WorkflowToIncidentExecution.__table__ +COMPLETED_STATUSES = ("success", "error", "providers_not_configured") + + +@dataclass(frozen=True) +class RetentionPolicy: + success_days: int = 30 + failure_days: int | None = 30 + success_log_days: int | None = None + keep_successes: int = 5 + batch_size: int = 100 + log_batch_size: int = 5000 + batch_pause_ms: int = 200 + max_runtime_seconds: int = 600 + statement_timeout_ms: int = 15000 + lock_timeout_ms: int = 1000 + + def __post_init__(self): + limits = { + "success_days": 36500, + "keep_successes": 100, + "batch_size": 1000, + "log_batch_size": 10000, + "max_runtime_seconds": 3600, + "statement_timeout_ms": 600000, + "lock_timeout_ms": 60000, + } + for name, maximum in limits.items(): + value = getattr(self, name) + if type(value) is not int or not 1 <= value <= maximum: + raise ValueError(f"{name} must be between 1 and {maximum}") + if ( + type(self.batch_pause_ms) is not int + or not 0 <= self.batch_pause_ms <= 10000 + ): + raise ValueError("batch_pause_ms must be between 0 and 10000") + for name in ("failure_days", "success_log_days"): + value = getattr(self, name) + if value is not None and ( + type(value) is not int or not 1 <= value <= 36500 + ): + raise ValueError(f"{name} must be a positive number of days") + if ( + self.success_log_days is not None + and self.success_log_days > self.success_days + ): + raise ValueError("success_log_days cannot exceed success_days") + + +class _BudgetExpired(Exception): + pass + + +def _timeouts(connection, policy, deadline): + remaining = int((deadline - time.monotonic()) * 1000) + if remaining <= 0: + raise _BudgetExpired + connection.execute( + select( + func.set_config( + "statement_timeout", + str(min(policy.statement_timeout_ms, remaining)), + True, + ) + ) + ) + connection.execute( + select( + func.set_config( + "lock_timeout", + str(min(policy.lock_timeout_ms, remaining)), + True, + ) + ) + ) + + +def _preflight(connection): + """Refuse unadapted dependencies or an unindexed foreign-key lookup.""" + dependencies = connection.execute( + text( + """ + SELECT child.relname, attribute.attname, constraint_info.convalidated, + constraint_info.conrelid IN ( + to_regclass('workflowexecutionlog'), + to_regclass('workflowtoalertexecution'), + to_regclass('workflowtoincidentexecution') + ) AS known_table + FROM pg_constraint constraint_info + JOIN pg_class child ON child.oid = constraint_info.conrelid + JOIN pg_attribute attribute ON attribute.attrelid = child.oid + AND attribute.attnum = ANY(constraint_info.conkey) + WHERE constraint_info.contype = 'f' + AND constraint_info.confrelid = to_regclass('workflowexecution') + """ + ) + ).all() + if any( + not row.known_table + or not row.convalidated + or row.attname != "workflow_execution_id" + for row in dependencies + ): + raise ValueError( + "Unadapted inbound workflowexecution dependency; inspect the schema" + ) + if {row.relname for row in dependencies} != { + LOGS.name, + ALERT_LINKS.name, + INCIDENT_LINKS.name, + }: + raise ValueError( + "Missing workflowexecution dependency constraints; inspect the schema" + ) + child_dependencies = connection.scalar( + text( + """ + SELECT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE contype = 'f' AND confrelid IN ( + to_regclass('workflowexecutionlog'), + to_regclass('workflowtoalertexecution'), + to_regclass('workflowtoincidentexecution') + ) + ) + """ + ) + ) + if child_dependencies: + raise ValueError( + "Unadapted inbound history-child dependency; inspect the schema" + ) + for table in (LOGS, ALERT_LINKS, INCIDENT_LINKS): + indexed = connection.scalar( + text( + """ + SELECT EXISTS ( + SELECT 1 FROM pg_index index_info + JOIN pg_attribute attribute + ON attribute.attrelid = index_info.indrelid + AND attribute.attnum = index_info.indkey[0] + WHERE index_info.indrelid = to_regclass(:table_name) + AND index_info.indisvalid AND index_info.indisready + AND index_info.indpred IS NULL + AND attribute.attname = 'workflow_execution_id' + ) + """ + ), + {"table_name": table.name}, + ) + if not indexed: + raise ValueError( + f"A valid non-partial index beginning with workflow_execution_id " + f"is required on {table.name}; prepare the index separately" + ) + + +def _protected_ids(connection, tenant_id, keep_successes): + success = ( + select( + EXECUTIONS.c.id, + func.row_number() + .over( + partition_by=EXECUTIONS.c.workflow_id, + order_by=(EXECUTIONS.c.started.desc(), EXECUTIONS.c.id.desc()), + ) + .label("position"), + ) + .where( + EXECUTIONS.c.tenant_id == tenant_id, + EXECUTIONS.c.status == "success", + EXECUTIONS.c.is_test_run.is_(False), + ) + .subquery() + ) + completed = ( + select( + EXECUTIONS.c.id, + func.row_number() + .over( + partition_by=EXECUTIONS.c.workflow_id, + order_by=( + EXECUTIONS.c.execution_number.desc(), + EXECUTIONS.c.started.desc(), + EXECUTIONS.c.id.desc(), + ), + ) + .label("position"), + ) + .where( + EXECUTIONS.c.tenant_id == tenant_id, + EXECUTIONS.c.status.in_(COMPLETED_STATUSES), + EXECUTIONS.c.is_test_run.is_(False), + ) + .subquery() + ) + protected = set( + connection.execute( + select(success.c.id).where(success.c.position <= keep_successes) + ).scalars() + ) + protected.update( + connection.execute( + select(completed.c.id).where(completed.c.position == 1) + ).scalars() + ) + return protected + + +def _old_completed(statuses, cutoff): + # The indexed started predicate is retained as a cheap preliminary filter. + # is_running is a scheduler sentinel/random integer, not a completion boolean. + return and_( + EXECUTIONS.c.status.in_(statuses), + EXECUTIONS.c.started < cutoff, + EXECUTIONS.c.execution_time >= 0, + EXECUTIONS.c.started + EXECUTIONS.c.execution_time * text("INTERVAL '1 second'") + < cutoff, + ) + + +def _predicates(policy, tenant_id, protected, now): + scope = and_(EXECUTIONS.c.tenant_id == tenant_id, EXECUTIONS.c.id.not_in(protected)) + success = _old_completed(("success",), now - timedelta(days=policy.success_days)) + terminal = [success] + if policy.failure_days is not None: + terminal.append( + _old_completed( + ("error", "providers_not_configured"), + now - timedelta(days=policy.failure_days), + ) + ) + parents = and_(scope, or_(*terminal)) + log_parents = parents + logs = parents + if policy.success_log_days is not None: + cutoff = now - timedelta(days=policy.success_log_days) + early_success = and_(scope, _old_completed(("success",), cutoff)) + log_parents = or_(parents, early_success) + logs = or_(parents, and_(early_success, LOGS.c.timestamp < cutoff)) + return parents, log_parents, logs + + +def _preview(connection, report, parents, logs, policy, deadline): + operations = [ + ( + "eligible_executions", + select(func.count()).select_from(EXECUTIONS).where(parents), + ), + ( + "stored_results_bytes", + select( + func.coalesce(func.sum(func.pg_column_size(EXECUTIONS.c.results)), 0) + ).where(parents), + ), + ( + "eligible_logs", + select(func.count()) + .select_from( + LOGS.join(EXECUTIONS, LOGS.c.workflow_execution_id == EXECUTIONS.c.id) + ) + .where(logs), + ), + ( + "eligible_alert_associations", + select(func.count()) + .select_from( + ALERT_LINKS.join( + EXECUTIONS, ALERT_LINKS.c.workflow_execution_id == EXECUTIONS.c.id + ) + ) + .where(parents), + ), + ( + "eligible_incident_associations", + select(func.count()) + .select_from( + INCIDENT_LINKS.join( + EXECUTIONS, + INCIDENT_LINKS.c.workflow_execution_id == EXECUTIONS.c.id, + ) + ) + .where(parents), + ), + ] + for key, statement in operations: + _timeouts(connection, policy, deadline) + value = connection.scalar(statement) + report[key] = (report[key] or 0) + value + + +def _delete_batch(connection, policy, parents, log_parents, logs): + """Never materialize results/log bodies; limit children separately.""" + has_logs = exists( + select(LOGS.c.id) + .where( + LOGS.c.workflow_execution_id == EXECUTIONS.c.id, + logs, + ) + .correlate(EXECUTIONS) + ) + log_parent_ids = list( + connection.execute( + select(EXECUTIONS.c.id) + .where(log_parents, has_logs) + .order_by(EXECUTIONS.c.started, EXECUTIONS.c.id) + .limit(policy.batch_size) + ).scalars() + ) + deleted = { + "deleted_logs": 0, + "deleted_executions": 0, + "deleted_alert_associations": 0, + "deleted_incident_associations": 0, + } + if log_parent_ids: + log_ids = ( + select(LOGS.c.id) + .select_from( + LOGS.join(EXECUTIONS, LOGS.c.workflow_execution_id == EXECUTIONS.c.id) + ) + .where( + EXECUTIONS.c.id.in_(log_parent_ids), + logs, + ) + .limit(policy.log_batch_size) + ) + deleted["deleted_logs"] = connection.execute( + delete(LOGS).where(LOGS.c.id.in_(log_ids)) + ).rowcount + no_logs = ~exists( + select(LOGS.c.id).where(LOGS.c.workflow_execution_id == EXECUTIONS.c.id) + ) + parent_ids = list( + connection.execute( + select(EXECUTIONS.c.id) + .where(parents, no_logs) + .order_by(EXECUTIONS.c.started, EXECUTIONS.c.id) + .limit(policy.batch_size) + ).scalars() + ) + if parent_ids: + for table, key in ( + (ALERT_LINKS, "deleted_alert_associations"), + (INCIDENT_LINKS, "deleted_incident_associations"), + ): + association_ids = ( + select(table.c.id) + .where(table.c.workflow_execution_id.in_(parent_ids)) + .limit(policy.log_batch_size) + ) + deleted[key] = connection.execute( + delete(table).where(table.c.id.in_(association_ids)) + ).rowcount + no_alert_links = ~exists( + select(ALERT_LINKS.c.id).where( + ALERT_LINKS.c.workflow_execution_id == EXECUTIONS.c.id + ) + ) + no_incident_links = ~exists( + select(INCIDENT_LINKS.c.id).where( + INCIDENT_LINKS.c.workflow_execution_id == EXECUTIONS.c.id + ) + ) + deleted["deleted_executions"] = connection.execute( + delete(EXECUTIONS).where( + EXECUTIONS.c.id.in_(parent_ids), + parents, + no_logs, + no_alert_links, + no_incident_links, + ) + ).rowcount + return deleted + + +def run_retention( + engine, policy, *, tenant_id=None, all_tenants=False, apply=False, now=None +): + """Run against PostgreSQL; dry-run is read-only and incomplete reports are explicit. + + No schema modification, checkpoint rewriting, alert deletion, or VACUUM is done. + Counts from apply mode describe committed work, not an estimate of total scope. + """ + if bool(tenant_id) == bool(all_tenants) or ( + tenant_id is not None and not tenant_id.strip() + ): + raise ValueError("Choose exactly one nonempty tenant_id or all_tenants") + if engine.dialect.name != "postgresql": + raise ValueError("Workflow retention currently requires PostgreSQL") + now = now or datetime.now(timezone.utc) + if now.tzinfo is not None: + now = now.astimezone(timezone.utc).replace(tzinfo=None) + deadline = time.monotonic() + policy.max_runtime_seconds + report = { + "mode": "apply" if apply else "dry-run", + "checked_at_utc": now.isoformat() + "Z", + "completed": False, + "stop_reason": None, + "tenants": 0, + "protected_executions": 0, + "eligible_executions": None, + "eligible_logs": None, + "stored_results_bytes": None, + "eligible_alert_associations": None, + "eligible_incident_associations": None, + "deleted_executions": 0, + "deleted_logs": 0, + "deleted_alert_associations": 0, + "deleted_incident_associations": 0, + "committed_batches": 0, + } + try: + with engine.connect() as connection: + # A preview uses one bounded, consistent read-only snapshot. + with connection.begin(): + connection.execute( + text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY") + ) + _timeouts(connection, policy, deadline) + _preflight(connection) + tenants = ( + [tenant_id] + if tenant_id + else list( + connection.execute( + select(EXECUTIONS.c.tenant_id) + .distinct() + .order_by(EXECUTIONS.c.tenant_id) + ).scalars() + ) + ) + protections = {} + for tenant in tenants: + _timeouts(connection, policy, deadline) + protections[tenant] = _protected_ids( + connection, tenant, policy.keep_successes + ) + report["protected_executions"] += len(protections[tenant]) + if not apply: + parents, _, logs = _predicates( + policy, tenant, protections[tenant], now + ) + _preview(connection, report, parents, logs, policy, deadline) + report["tenants"] += 1 + if apply: + for tenant in tenants: + parents, log_parents, logs = _predicates( + policy, tenant, protections[tenant], now + ) + while True: + with connection.begin(): + _timeouts(connection, policy, deadline) + deleted = _delete_batch( + connection, policy, parents, log_parents, logs + ) + # Only count the batch after its transaction commits. + if not any(deleted.values()): + break + for key, value in deleted.items(): + report[key] += value + report["committed_batches"] += 1 + logger.info("Committed workflow retention batch: %s", deleted) + pause = min( + policy.batch_pause_ms / 1000, + max(0, deadline - time.monotonic()), + ) + if pause: + time.sleep(pause) + if not apply: + for key in ( + "eligible_executions", + "eligible_logs", + "stored_results_bytes", + "eligible_alert_associations", + "eligible_incident_associations", + ): + report[key] = report[key] or 0 + report["completed"] = True + report["stop_reason"] = "finished" + except _BudgetExpired: + report["stop_reason"] = "runtime_limit" + except DBAPIError as error: + code = getattr(error.orig, "sqlstate", None) or getattr( + error.orig, "pgcode", None + ) + if code != "57014": + raise + report["stop_reason"] = "query_timeout" + return report diff --git a/keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py b/keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py new file mode 100644 index 0000000000..b5a64bbd4a --- /dev/null +++ b/keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py @@ -0,0 +1,86 @@ +"""Index workflow execution log lookups for bounded history maintenance. + +Revision ID: 6942a1b2c3d4 +Revises: 67ff7efffed4 +Create Date: 2026-11-12 13:30:00 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "6942a1b2c3d4" +down_revision = "67ff7efffed4" +branch_labels = None +depends_on = None + +INDEX = "idx_workflowexecutionlog_workflow_execution_id" +TABLE = "workflowexecutionlog" +COLUMN = "workflow_execution_id" + + +def upgrade(): + connection = op.get_bind() + if connection.dialect.name == "postgresql": + indexes = ( + connection.execute( + sa.text( + """ + SELECT index_relation.relname AS name, index_info.indisvalid AS valid, + index_info.indisready AS ready, index_info.indpred IS NULL AS full_index, + attribute.attname AS first_column + FROM pg_index index_info + JOIN pg_class index_relation ON index_relation.oid = index_info.indexrelid + LEFT JOIN pg_attribute attribute ON attribute.attrelid = index_info.indrelid + AND attribute.attnum = index_info.indkey[0] + WHERE index_info.indrelid = to_regclass('workflowexecutionlog') + """ + ) + ) + .mappings() + .all() + ) + for index in indexes: + if index["name"] == INDEX and not ( + index["valid"] + and index["ready"] + and index["full_index"] + and index["first_column"] == COLUMN + ): + raise RuntimeError( + f"Inspect/repair {INDEX} before retrying: it is invalid or has a different definition" + ) + if any( + index["valid"] + and index["ready"] + and index["full_index"] + and index["first_column"] == COLUMN + for index in indexes + ): + return + # Existing large installations may prebuild this index before upgrading. + # Never wrap CREATE INDEX CONCURRENTLY in the migration transaction. + with op.get_context().autocommit_block(): + op.create_index(INDEX, TABLE, [COLUMN], postgresql_concurrently=True) + else: + indexes = sa.inspect(connection).get_indexes(TABLE) + if any( + index["column_names"] + and index["column_names"][0] == COLUMN + and index.get("dialect_options", {}).get("sqlite_where") is None + for index in indexes + ): + return + op.create_index(INDEX, TABLE, [COLUMN]) + + +def downgrade(): + connection = op.get_bind() + if not any( + index["name"] == INDEX for index in sa.inspect(connection).get_indexes(TABLE) + ): + return + if connection.dialect.name == "postgresql": + with op.get_context().autocommit_block(): + op.drop_index(INDEX, table_name=TABLE, postgresql_concurrently=True) + else: + op.drop_index(INDEX, table_name=TABLE) diff --git a/keep/api/models/db/workflow.py b/keep/api/models/db/workflow.py index 900fb6bed3..d2f154f6a3 100644 --- a/keep/api/models/db/workflow.py +++ b/keep/api/models/db/workflow.py @@ -195,6 +195,12 @@ class WorkflowToIncidentExecution(SQLModel, table=True): class WorkflowExecutionLog(SQLModel, table=True): + __table_args__ = ( + Index( + "idx_workflowexecutionlog_workflow_execution_id", "workflow_execution_id" + ), + ) + id: int = Field(default=None, primary_key=True) workflow_execution_id: str = Field(foreign_key="workflowexecution.id") timestamp: datetime diff --git a/keep/cli/workflow_retention.py b/keep/cli/workflow_retention.py new file mode 100644 index 0000000000..0128d01d0d --- /dev/null +++ b/keep/cli/workflow_retention.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Preview/apply bounded PostgreSQL workflow-history retention (issue #6942).""" + +import argparse +import json +import logging +import sys + +from keep.api.core.workflow_retention import RetentionPolicy, run_retention + + +def _get_engine(): + # Reuse Keep's configured engine without treating get_session() as a context manager. + # Import only after validating arguments; --help/invalid arguments never open the DB. + from keep.api.core.db import engine + + return engine + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + scope = parser.add_mutually_exclusive_group(required=True) + scope.add_argument("--tenant-id") + scope.add_argument("--all-tenants", action="store_true") + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--apply", + action="store_true", + help="Delete eligible history; default is dry-run", + ) + mode.add_argument( + "--dry-run", + action="store_true", + help="Explicit read-only preview (the default)", + ) + parser.add_argument( + "--retention-days", + "--success-retention-days", + dest="success_days", + type=int, + default=30, + ) + failures = parser.add_mutually_exclusive_group() + failures.add_argument( + "--failure-retention-days", dest="failure_days", type=int, default=30 + ) + failures.add_argument("--keep-failed", action="store_true") + parser.add_argument( + "--success-log-retention-days", dest="success_log_days", type=int + ) + parser.add_argument("--keep-successes", type=int, default=5) + parser.add_argument("--batch-size", type=int, default=100) + parser.add_argument("--log-batch-size", type=int, default=5000) + parser.add_argument("--batch-pause-ms", type=int, default=200) + parser.add_argument("--max-runtime-seconds", type=int, default=600) + parser.add_argument("--statement-timeout-ms", type=int, default=15000) + parser.add_argument("--lock-timeout-ms", type=int, default=1000) + args = parser.parse_args(argv) + try: + policy = RetentionPolicy( + success_days=args.success_days, + failure_days=None if args.keep_failed else args.failure_days, + success_log_days=args.success_log_days, + keep_successes=args.keep_successes, + batch_size=args.batch_size, + log_batch_size=args.log_batch_size, + batch_pause_ms=args.batch_pause_ms, + max_runtime_seconds=args.max_runtime_seconds, + statement_timeout_ms=args.statement_timeout_ms, + lock_timeout_ms=args.lock_timeout_ms, + ) + if args.tenant_id is not None and not args.tenant_id.strip(): + raise ValueError("tenant-id cannot be empty") + except ValueError as error: + parser.error(str(error)) + try: + report = run_retention( + _get_engine(), + policy, + tenant_id=args.tenant_id, + all_tenants=args.all_tenants, + apply=args.apply, + ) + except ValueError as error: + parser.exit(1, f"Retention preflight failed: {error}\n") + except Exception as error: + # Driver exceptions can embed connection credentials, SQL, or row contents. + parser.exit( + 1, + f"Retention failed ({type(error).__name__}); check database health and permissions\n", + ) + print(json.dumps(report, sort_keys=True)) + return 0 if report["completed"] else 2 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + sys.exit(main()) diff --git a/scripts/cleanup_workflow_executions.py b/scripts/cleanup_workflow_executions.py new file mode 100644 index 0000000000..3207188238 --- /dev/null +++ b/scripts/cleanup_workflow_executions.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +"""Source-checkout entry point for PostgreSQL workflow-history maintenance.""" + +import sys +from pathlib import Path + +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from keep.cli.workflow_retention import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/maintenance/conftest.py b/tests/maintenance/conftest.py new file mode 100644 index 0000000000..6939c1dd90 --- /dev/null +++ b/tests/maintenance/conftest.py @@ -0,0 +1,132 @@ +"""Isolated PostgreSQL tests; never use Keep's application connection string.""" + +import os +from datetime import datetime, timedelta +from uuid import uuid4 + +import pytest +from sqlalchemy import create_engine, text +from sqlmodel import SQLModel + +from keep.api.models.db.tenant import Tenant +from keep.api.models.db.workflow import ( + Workflow, + WorkflowExecution, + WorkflowExecutionLog, + WorkflowToAlertExecution, + WorkflowToIncidentExecution, + WorkflowVersion, +) + +NOW = datetime(2026, 11, 12, 12) + + +@pytest.fixture +def pg_engine(): + url = os.getenv("KEEP_RETENTION_TEST_DATABASE_URL") + if not url: + pytest.skip("Set KEEP_RETENTION_TEST_DATABASE_URL to an isolated PostgreSQL DB") + admin = create_engine(url) + schema = "keep_retention_test_" + uuid4().hex + with admin.begin() as connection: + connection.execute(text(f'CREATE SCHEMA "{schema}"')) + engine = create_engine( + url, connect_args={"options": f"-c search_path={schema},public"} + ) + tables = [ + Tenant.__table__, + Workflow.__table__, + WorkflowVersion.__table__, + WorkflowExecution.__table__, + WorkflowExecutionLog.__table__, + WorkflowToAlertExecution.__table__, + WorkflowToIncidentExecution.__table__, + ] + SQLModel.metadata.create_all(engine, tables=tables) + try: + yield engine + finally: + engine.dispose() + with admin.begin() as connection: + connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + admin.dispose() + + +@pytest.fixture +def seed(pg_engine): + tenants, workflows = set(), set() + sequence = 0 + + def add( + identifier, + *, + days=45, + status="success", + tenant="tenant-a", + workflow="wf-a", + is_test=False, + execution_time=1, + logs=0, + associations=False, + result=None, + ): + nonlocal sequence + sequence += 1 + with pg_engine.begin() as connection: + if tenant not in tenants: + connection.execute( + Tenant.__table__.insert().values(id=tenant, name=tenant) + ) + tenants.add(tenant) + if workflow not in workflows: + connection.execute( + Workflow.__table__.insert().values( + id=workflow, + tenant_id=tenant, + name=workflow, + created_by="retention-test", + workflow_raw="{}", + interval=60, + ) + ) + workflows.add(workflow) + connection.execute( + WorkflowExecution.__table__.insert().values( + id=identifier, + tenant_id=tenant, + workflow_id=workflow, + status=status, + started=NOW - timedelta(days=days), + triggered_by="scheduler", + execution_number=sequence, + execution_time=execution_time, + is_test_run=is_test, + is_running=12345 if status != "in_progress" else 1, + results=result or {"checkpoint": identifier}, + ) + ) + for number in range(logs): + connection.execute( + WorkflowExecutionLog.__table__.insert().values( + workflow_execution_id=identifier, + timestamp=NOW - timedelta(days=days), + message=f"test log {number}", + context={}, + ) + ) + if associations: + connection.execute( + WorkflowToAlertExecution.__table__.insert().values( + workflow_execution_id=identifier, + alert_fingerprint="alert-1", + ) + ) + connection.execute( + WorkflowToIncidentExecution.__table__.insert().values( + workflow_execution_id=identifier, + incident_id="incident-1", + ) + ) + return identifier + + return add diff --git a/tests/maintenance/test_workflow_retention.py b/tests/maintenance/test_workflow_retention.py new file mode 100644 index 0000000000..6cb47224eb --- /dev/null +++ b/tests/maintenance/test_workflow_retention.py @@ -0,0 +1,517 @@ +"""Retention correctness against real PostgreSQL constraints and transactions.""" + +import importlib.util +import subprocess +import sys +from datetime import datetime +from pathlib import Path + +import pytest +from sqlalchemy import event, func, select, text + +from keep.api.core.workflow_retention import RetentionPolicy, run_retention +from keep.api.models.db.workflow import ( + WorkflowExecution, + WorkflowExecutionLog, + WorkflowToAlertExecution, + WorkflowToIncidentExecution, +) + +NOW = datetime(2026, 11, 12, 12) + +EXECUTIONS = WorkflowExecution.__table__ +LOGS = WorkflowExecutionLog.__table__ +ALERT_LINKS = WorkflowToAlertExecution.__table__ +INCIDENT_LINKS = WorkflowToIncidentExecution.__table__ +INDEX = "idx_workflowexecutionlog_workflow_execution_id" + + +def ids(engine): + with engine.connect() as connection: + return set(connection.execute(select(EXECUTIONS.c.id)).scalars()) + + +def count(engine, table): + with engine.connect() as connection: + return connection.scalar(select(func.count()).select_from(table)) + + +@pytest.mark.parametrize( + "field,value", + [ + ("success_days", -1), + ("success_days", 0), + ("failure_days", 0), + ("batch_size", 0), + ("batch_size", 1001), + ("log_batch_size", 0), + ("log_batch_size", 10001), + ("keep_successes", 0), + ("batch_pause_ms", -1), + ("max_runtime_seconds", 0), + ("statement_timeout_ms", 0), + ("success_log_days", 31), + ], +) +def test_reject_invalid_policy(field, value): + with pytest.raises(ValueError): + RetentionPolicy(**{field: value}) + + +def test_scope_is_required(pg_engine): + with pytest.raises(ValueError): + run_retention(pg_engine, RetentionPolicy(), now=NOW) + with pytest.raises(ValueError): + run_retention( + pg_engine, RetentionPolicy(), tenant_id="a", all_tenants=True, now=NOW + ) + + +def test_dry_run_counts_all_batches_without_mutations(pg_engine, seed): + for number in range(5): + seed(f"old-{number}", logs=2, associations=True) + seed("recent", days=1) + statements = [] + + def record(_, __, statement, *args): + statements.append(statement) + + event.listen(pg_engine, "before_cursor_execute", record) + report = run_retention( + pg_engine, + RetentionPolicy(batch_size=2, keep_successes=1), + tenant_id="tenant-a", + now=NOW, + ) + assert report["completed"] + assert report["eligible_executions"] == 5 + assert report["eligible_logs"] == 10 + assert report["eligible_alert_associations"] == 5 + assert report["eligible_incident_associations"] == 5 + assert count(pg_engine, EXECUTIONS) == 6 + assert count(pg_engine, LOGS) == 10 + assert not any( + s.lstrip().upper().startswith(("DELETE", "UPDATE", "INSERT")) + for s in statements + ) + + +def test_apply_removes_children_preserves_recent_active_and_other_tenant( + pg_engine, seed +): + seed("old", logs=7, associations=True) + seed("failed", days=45, status="error", logs=1, associations=True) + seed("recent", days=1) + seed("active", days=60, status="in_progress", execution_time=None, logs=1) + seed("timeout", days=60, status="timeout", execution_time=None) + seed("ambiguous", days=60, execution_time=None) + seed("other", tenant="tenant-b", workflow="wf-b", logs=1) + report = run_retention( + pg_engine, + RetentionPolicy(batch_size=1, log_batch_size=2, keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert report["completed"] + assert report["deleted_executions"] == 2 + assert report["deleted_logs"] == 8 + assert report["deleted_alert_associations"] == 2 + assert report["deleted_incident_associations"] == 2 + assert ids(pg_engine) == {"recent", "active", "timeout", "ambiguous", "other"} + assert count(pg_engine, LOGS) == 2 + again = run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert again["deleted_executions"] == 0 + + +def test_protect_latest_successes_and_scheduler_record_even_when_old(pg_engine, seed): + seed("old-success-1", days=100) + seed("old-success-2", days=90) + seed("test-success", days=80, is_test=True) + seed("last-completed", days=70, status="error") + report = run_retention( + pg_engine, + RetentionPolicy(keep_successes=2), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert report["deleted_executions"] == 1 + assert ids(pg_engine) == {"old-success-1", "old-success-2", "last-completed"} + + +def test_failure_retention_and_real_completion_time(pg_engine, seed): + seed("old", days=20) + seed("recent-failure", days=20, status="error") + seed("old-failure", days=40, status="error") + seed("long-running", days=40, execution_time=39 * 86400) + seed("recent", days=1) + run_retention( + pg_engine, + RetentionPolicy(success_days=14, failure_days=30, keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert ids(pg_engine) == {"recent-failure", "long-running", "recent"} + + +def test_success_log_retention_keeps_parent_and_failure_logs(pg_engine, seed): + seed("success", days=10, logs=7) + seed("error", days=10, status="error", logs=3) + seed("recent", days=1) + report = run_retention( + pg_engine, + RetentionPolicy( + success_days=14, success_log_days=7, log_batch_size=2, keep_successes=1 + ), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert report["deleted_executions"] == 0 + assert report["deleted_logs"] == 7 + assert count(pg_engine, EXECUTIONS) == 3 + assert count(pg_engine, LOGS) == 3 + + +def test_missing_log_index_refuses_apply(pg_engine, seed): + seed("old", logs=1) + seed("recent", days=1) + with pg_engine.begin() as connection: + connection.execute(text(f'DROP INDEX "{INDEX}"')) + with pytest.raises(ValueError, match="index"): + run_retention( + pg_engine, RetentionPolicy(), tenant_id="tenant-a", apply=True, now=NOW + ) + assert count(pg_engine, EXECUTIONS) == 2 + assert count(pg_engine, LOGS) == 1 + + +def test_unknown_inbound_dependency_refuses_cleanup(pg_engine, seed): + seed("old") + with pg_engine.begin() as connection: + connection.execute( + text( + "CREATE TABLE future_dependency (id integer PRIMARY KEY, execution_id varchar REFERENCES workflowexecution(id))" + ) + ) + with pytest.raises(ValueError, match="depend"): + run_retention( + pg_engine, RetentionPolicy(), tenant_id="tenant-a", apply=True, now=NOW + ) + + +def test_rollback_does_not_delete_parent_or_associations(pg_engine, seed): + seed("old", associations=True) + seed("recent", days=1) + + def fail(_, __, statement, *args): + if statement.startswith("DELETE FROM workflowexecution "): + raise RuntimeError("injected failure before parent deletion") + + event.listen(pg_engine, "before_cursor_execute", fail) + with pytest.raises(RuntimeError, match="injected failure"): + run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + event.remove(pg_engine, "before_cursor_execute", fail) + assert count(pg_engine, EXECUTIONS) == 2 + assert count(pg_engine, ALERT_LINKS) == 1 + assert count(pg_engine, INCIDENT_LINKS) == 1 + + +def test_all_tenants_are_processed_with_separate_protection(pg_engine, seed): + for tenant, workflow in [("tenant-a", "wf-a"), ("tenant-b", "wf-b")]: + seed(tenant + "-old", tenant=tenant, workflow=workflow) + seed(tenant + "-recent", tenant=tenant, workflow=workflow, days=1) + report = run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + all_tenants=True, + apply=True, + now=NOW, + ) + assert report["deleted_executions"] == 2 + assert ids(pg_engine) == {"tenant-a-recent", "tenant-b-recent"} + + +def test_cli_invalid_arguments_and_help_do_not_load_application_db(monkeypatch): + path = ( + Path(__file__).resolve().parents[2] / "keep" / "cli" / "workflow_retention.py" + ) + spec = importlib.util.spec_from_file_location("retention_cli_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + for arguments, code in [ + (["--help"], 0), + ([], 2), + (["--tenant-id", "tenant-a", "--retention-days", "-1"], 2), + ]: + with pytest.raises(SystemExit) as error: + module.main(arguments) + assert error.value.code == code + + +def test_association_fanout_is_bounded(pg_engine, seed): + seed("old") + seed("recent", days=1) + with pg_engine.begin() as connection: + connection.execute( + ALERT_LINKS.insert(), + [ + dict(workflow_execution_id="old", alert_fingerprint=f"a-{n}") + for n in range(9) + ], + ) + report = run_retention( + pg_engine, + RetentionPolicy(log_batch_size=2, keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert report["deleted_alert_associations"] == 9 + assert report["deleted_executions"] == 1 + assert report["committed_batches"] == 5 + + +def test_runtime_limit_reports_committed_progress_and_can_resume( + pg_engine, seed, monkeypatch +): + from keep.api.core import workflow_retention + + for number in range(3): + seed(f"old-{number}") + seed("recent", days=1) + clock = [0.0] + monkeypatch.setattr(workflow_retention.time, "monotonic", lambda: clock[0]) + + def advance(_, __, statement, *args): + if statement.startswith("DELETE FROM workflowexecution "): + clock[0] += 2 + + event.listen(pg_engine, "after_cursor_execute", advance) + report = run_retention( + pg_engine, + RetentionPolicy(batch_size=1, keep_successes=1, max_runtime_seconds=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + event.remove(pg_engine, "after_cursor_execute", advance) + assert not report["completed"] + assert report["stop_reason"] == "runtime_limit" + assert report["deleted_executions"] == 1 + assert count(pg_engine, EXECUTIONS) == 3 + resumed = run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert resumed["deleted_executions"] == 2 + + +def test_query_timeout_marks_preview_incomplete(pg_engine, seed): + seed("old") + seed("recent", days=1) + + def slow_count(connection, cursor, statement, parameters, context, executemany): + if statement.startswith("SELECT count("): + return "SELECT pg_sleep(0.1)", () + return statement, parameters + + event.listen(pg_engine, "before_cursor_execute", slow_count, retval=True) + report = run_retention( + pg_engine, + RetentionPolicy(keep_successes=1, statement_timeout_ms=10), + tenant_id="tenant-a", + now=NOW, + ) + event.remove(pg_engine, "before_cursor_execute", slow_count) + assert not report["completed"] + assert report["stop_reason"] == "query_timeout" + assert report["deleted_executions"] == 0 + assert count(pg_engine, EXECUTIONS) == 2 + + +def test_apply_never_selects_result_payload(pg_engine, seed): + seed("old", result={"large": "x" * 100000}) + seed("recent", days=1) + statements = [] + + def record(_, __, statement, *args): + statements.append(statement) + + event.listen(pg_engine, "before_cursor_execute", record) + run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert not any("workflowexecution.results" in s for s in statements) + + +def test_terminal_provider_configuration_error_is_eligible(pg_engine, seed): + seed("old", status="providers_not_configured") + seed("recent", days=1) + report = run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert report["deleted_executions"] == 1 + + +def test_cli_success_uses_synchronous_engine(pg_engine, seed, monkeypatch, capsys): + import json + + path = ( + Path(__file__).resolve().parents[2] / "keep" / "cli" / "workflow_retention.py" + ) + spec = importlib.util.spec_from_file_location("retention_cli_success", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + seed("old", days=365) + seed("recent", days=0) + monkeypatch.setattr(module, "_get_engine", lambda: pg_engine) + assert module.main(["--tenant-id", "tenant-a", "--keep-successes", "1"]) == 0 + report = json.loads(capsys.readouterr().out) + assert report["mode"] == "dry-run" + assert report["eligible_executions"] == 1 + assert count(pg_engine, EXECUTIONS) == 2 + + +@pytest.mark.parametrize( + "entrypoint", + [ + ["-m", "keep.cli.workflow_retention"], + ["scripts/cleanup_workflow_executions.py"], + ], +) +def test_packaged_and_source_cli_help(entrypoint): + result = subprocess.run( + [sys.executable, *entrypoint, "--help"], + capture_output=True, + text=True, + timeout=10, + ) + assert result.returncode == 0 + assert "--apply" in result.stdout + + +def test_zero_candidates_preview_is_complete(pg_engine): + report = run_retention(pg_engine, RetentionPolicy(), all_tenants=True, now=NOW) + assert report["completed"] + assert report["eligible_executions"] == 0 + assert report["eligible_logs"] == 0 + + +def test_keep_failed_preserves_old_errors(pg_engine, seed): + seed("error", status="error", logs=2) + seed("recent", days=1) + run_retention( + pg_engine, + RetentionPolicy(failure_days=None, keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert ids(pg_engine) == {"error", "recent"} + assert count(pg_engine, LOGS) == 2 + + +def test_protected_checkpoint_is_not_rewritten(pg_engine, seed): + payload = { + "sync-slow-query-patterns": [{"checkpoint_base64": "b3BhcXVl", "delivered": 0}], + "extra": [1, 2, 3], + } + seed("protected", days=100, logs=2, result=payload) + seed("old-test", days=90, is_test=True) + run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + with pg_engine.connect() as connection: + assert ( + connection.scalar( + select(EXECUTIONS.c.results).where(EXECUTIONS.c.id == "protected") + ) + == payload + ) + assert count(pg_engine, LOGS) == 2 + + +def test_missing_fk_refuses_apply(pg_engine): + with pg_engine.begin() as connection: + connection.execute( + text( + "ALTER TABLE workflowexecutionlog DROP CONSTRAINT workflowexecutionlog_workflow_execution_id_fkey" + ) + ) + with pytest.raises(ValueError, match="depend"): + run_retention( + pg_engine, RetentionPolicy(), tenant_id="tenant-a", apply=True, now=NOW + ) + + +def test_partial_log_index_is_not_accepted(pg_engine, seed): + seed("old") + with pg_engine.begin() as connection: + connection.execute(text(f'DROP INDEX "{INDEX}"')) + connection.execute( + text( + f'CREATE INDEX "{INDEX}" ON workflowexecutionlog(workflow_execution_id) WHERE workflow_execution_id IS NOT NULL' + ) + ) + with pytest.raises(ValueError, match="index"): + run_retention( + pg_engine, RetentionPolicy(), tenant_id="tenant-a", apply=True, now=NOW + ) + + +@pytest.mark.parametrize( + "parent", + ["workflowexecutionlog", "workflowtoalertexecution", "workflowtoincidentexecution"], +) +def test_unknown_history_child_dependency_refuses_cascading_delete( + pg_engine, seed, parent +): + seed("old", logs=1, associations=True) + seed("recent", days=1) + with pg_engine.begin() as connection: + connection.execute( + text( + f"CREATE TABLE future_child_dependency (id integer PRIMARY KEY, history_id integer REFERENCES {parent}(id) ON DELETE CASCADE)" + ) + ) + with pytest.raises(ValueError, match="depend"): + run_retention( + pg_engine, + RetentionPolicy(keep_successes=1), + tenant_id="tenant-a", + apply=True, + now=NOW, + ) + assert count(pg_engine, LOGS) == 1 + assert count(pg_engine, ALERT_LINKS) == 1 + assert count(pg_engine, INCIDENT_LINKS) == 1 diff --git a/tests/maintenance/test_workflow_retention_migration.py b/tests/maintenance/test_workflow_retention_migration.py new file mode 100644 index 0000000000..bc49670f09 --- /dev/null +++ b/tests/maintenance/test_workflow_retention_migration.py @@ -0,0 +1,127 @@ +"""Exercise the actual Alembic index revision outside application startup.""" + +import importlib.util +from pathlib import Path + +import pytest +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import inspect, text +from sqlalchemy.exc import DBAPIError + +PATH = ( + Path(__file__).resolve().parents[2] + / "keep/api/models/db/migrations/versions/2026-11-12-13-30_6942a1b2c3d4.py" +) +INDEX = "idx_workflowexecutionlog_workflow_execution_id" + + +def revision(): + spec = importlib.util.spec_from_file_location("workflow_retention_index", PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def migrate(engine, direction): + with engine.connect() as connection: + context = MigrationContext.configure(connection) + with context.begin_transaction(): + with Operations.context(context): + getattr(revision(), direction)() + + +def test_concurrent_index_upgrade_is_idempotent_and_downgrade_works(pg_engine): + with pg_engine.begin() as connection: + connection.execute(text(f'DROP INDEX "{INDEX}"')) + migrate(pg_engine, "upgrade") + migrate(pg_engine, "upgrade") + assert INDEX in { + i["name"] for i in inspect(pg_engine).get_indexes("workflowexecutionlog") + } + migrate(pg_engine, "downgrade") + assert INDEX not in { + i["name"] for i in inspect(pg_engine).get_indexes("workflowexecutionlog") + } + + +def test_invalid_index_requires_explicit_repair(pg_engine, seed): + seed("old", logs=2) + with pg_engine.begin() as connection: + connection.execute(text(f'DROP INDEX "{INDEX}"')) + with pg_engine.connect().execution_options( + isolation_level="AUTOCOMMIT" + ) as connection: + with pytest.raises(DBAPIError): + connection.execute( + text( + f'CREATE UNIQUE INDEX CONCURRENTLY "{INDEX}" ON workflowexecutionlog(workflow_execution_id)' + ) + ) + with pytest.raises(RuntimeError, match="Inspect/repair"): + migrate(pg_engine, "upgrade") + + +def test_accepts_prebuilt_equivalent_index(pg_engine): + with pg_engine.begin() as connection: + connection.execute(text(f'DROP INDEX "{INDEX}"')) + connection.execute( + text( + "CREATE INDEX prebuilt_execution_lookup ON workflowexecutionlog(workflow_execution_id)" + ) + ) + migrate(pg_engine, "upgrade") + names = {i["name"] for i in inspect(pg_engine).get_indexes("workflowexecutionlog")} + assert "prebuilt_execution_lookup" in names + assert INDEX not in names + + +def test_revision_has_one_head(): + # Validate the revision graph without importing unrelated historical migration + # modules (which require the entire application's optional provider stack). + import ast + + root = Path(__file__).resolve().parents[2] + revisions, parents = set(), set() + for path in (root / "keep/api/models/db/migrations/versions").glob("*.py"): + metadata = {} + for node in ast.parse(path.read_text()).body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id in ( + "revision", + "down_revision", + ): + metadata[target.id] = ast.literal_eval(node.value) + if "revision" in metadata: + revisions.add(metadata["revision"]) + parent = metadata.get("down_revision") + parents.update( + parent if isinstance(parent, tuple) else [parent] if parent else [] + ) + assert parents <= revisions + assert revisions - parents == {"6942a1b2c3d4"} + + +def test_sqlite_migration_upgrade_and_downgrade_remain_compatible(tmp_path): + from sqlalchemy import create_engine + + engine = create_engine(f"sqlite:///{tmp_path / 'migration.db'}") + try: + with engine.begin() as connection: + connection.execute( + text( + "CREATE TABLE workflowexecutionlog (id integer PRIMARY KEY, workflow_execution_id varchar NOT NULL)" + ) + ) + migrate(engine, "upgrade") + migrate(engine, "upgrade") + assert INDEX in { + i["name"] for i in inspect(engine).get_indexes("workflowexecutionlog") + } + migrate(engine, "downgrade") + assert INDEX not in { + i["name"] for i in inspect(engine).get_indexes("workflowexecutionlog") + } + finally: + engine.dispose()