From 93ab9c786392fa0e3d08a45d75f4999087178c62 Mon Sep 17 00:00:00 2001 From: granbellllc-png Date: Wed, 11 Nov 2026 15:18:30 +0400 Subject: [PATCH 1/3] fix(servicenow): safely parse result list and annotate return type in validate_scopes (#6957) --- .../servicenow_provider.py | 10 ++++- tests/test_servicenow_provider.py | 37 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/keep/providers/servicenow_provider/servicenow_provider.py b/keep/providers/servicenow_provider/servicenow_provider.py index 3c674ac8b1..007a01b6b1 100644 --- a/keep/providers/servicenow_provider/servicenow_provider.py +++ b/keep/providers/servicenow_provider/servicenow_provider.py @@ -200,7 +200,7 @@ def _get_headers(self): headers["Authorization"] = f"Bearer {self._access_token}" return headers - def validate_scopes(self): + def validate_scopes(self) -> dict[str, bool | str]: """ Validates that the user has the required scopes to use the provider. """ @@ -244,7 +244,13 @@ def validate_scopes(self): if response.ok: roles = response.json() - roles_names = [role.get("name") for role in roles.get("result")] + raw_result = roles.get("result") if isinstance(roles, dict) else [] + result_list = raw_result if isinstance(raw_result, list) else [] + roles_names = [ + role.get("name") + for role in result_list + if isinstance(role, dict) + ] if "itil" in roles_names: self.logger.info("User has ITIL role") scopes = { diff --git a/tests/test_servicenow_provider.py b/tests/test_servicenow_provider.py index 90be793566..d2c3e364f3 100644 --- a/tests/test_servicenow_provider.py +++ b/tests/test_servicenow_provider.py @@ -481,3 +481,40 @@ def test_cmdb_fields_are_separate_entries(self, servicenow_provider): assert "owned_by.name" in fields assert "manufacturer.name" in fields assert "owned_by.namemanufacturer.name" not in fields + + +class TestValidateScopes: + """Tests for validate_scopes error handling and safety.""" + + def test_validate_scopes_with_itil_role(self, servicenow_provider): + mock_response = MagicMock() + mock_response.ok = True + mock_response.json.return_value = {"result": [{"name": "itil"}]} + + with patch("requests.get", return_value=mock_response): + scopes = servicenow_provider.validate_scopes() + assert scopes == {"itil": True} + + def test_validate_scopes_without_itil_role(self, servicenow_provider): + mock_response = MagicMock() + mock_response.ok = True + mock_response.json.return_value = {"result": [{"name": "admin"}]} + + with patch("requests.get", return_value=mock_response): + scopes = servicenow_provider.validate_scopes() + assert scopes == {"itil": "This user does not have the ITIL role"} + + def test_validate_scopes_missing_result_key(self, servicenow_provider): + """ServiceNow returns dict without 'result' key.""" + mock_response = MagicMock() + mock_response.ok = True + mock_response.json.return_value = {} + + with patch("requests.get", return_value=mock_response): + scopes = servicenow_provider.validate_scopes() + assert scopes == {"itil": "This user does not have the ITIL role"} + + def test_validate_scopes_skip_env(self, servicenow_provider, monkeypatch): + monkeypatch.setenv("KEEP_SERVICENOW_PROVIDER_SKIP_SCOPE_VALIDATION", "true") + scopes = servicenow_provider.validate_scopes() + assert scopes == {"itil": True} From 528d9f48e8e5178625a00faab56256fb412f89f5 Mon Sep 17 00:00:00 2001 From: granbellllc-png Date: Wed, 11 Nov 2026 20:44:26 +0400 Subject: [PATCH 2/3] feat(retention): add workflow execution and log cleanup maintenance script (#6942) Addresses #6942 by providing a supported database maintenance script to purge historical workflow executions and logs: - Bounded batch deletion chunks to prevent lock contention and WAL growth - Configurable retention window (--retention-days) and batch size (--batch-size) - Safe dry-run mode (--dry-run) to inspect purge counts without committing - Tenant isolation (--tenant-id) and error preservation (--keep-failed) - Fully covered by unit tests verifying isolation of active/fresh runs --- scripts/cleanup_workflow_executions.py | 132 ++++++++++++++++++++++ tests/test_cleanup_workflow_executions.py | 124 ++++++++++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 scripts/cleanup_workflow_executions.py create mode 100644 tests/test_cleanup_workflow_executions.py diff --git a/scripts/cleanup_workflow_executions.py b/scripts/cleanup_workflow_executions.py new file mode 100644 index 0000000000..866a25d3f1 --- /dev/null +++ b/scripts/cleanup_workflow_executions.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +""" +scripts/cleanup_workflow_executions.py + +PostgreSQL (and MySQL/SQLite) supported workflow execution and log retention cleanup. +Addresses #6942 (PostgreSQL TOAST/heap storage exhaustion on WorkflowExecution & WorkflowExecutionLog). + +Supports: +- --retention-days: purge completed executions older than N days (default: 30) +- --batch-size: purge in bounded chunks to avoid lock contention & WAL spikes (default: 500) +- --dry-run: simulate count of executions and logs that would be purged without deleting +- --tenant-id: limit purge to a specific tenant (default: all tenants) +- --keep-failed: preserve failed/error executions for post-mortem debugging +""" + +import argparse +import logging +import os +import sys +from datetime import datetime, timedelta, timezone + +from sqlmodel import Session, select, col + +# Adjust path for keep imports +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from keep.api.core.db import get_session +from keep.api.models.db.workflow import WorkflowExecution, WorkflowExecutionLog + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger(__name__) + +TERMINAL_SUCCESS_STATUSES = ["success", "passed"] +TERMINAL_FAILED_STATUSES = ["error", "timeout", "failed", "cancelled"] + + +def purge_workflow_executions( + session: Session, + retention_days: int = 30, + batch_size: int = 500, + dry_run: bool = False, + tenant_id: str | None = None, + keep_failed: bool = False, +) -> dict[str, int]: + cutoff_date = datetime.now(timezone.utc) - timedelta(days=retention_days) + logger.info(f"Purge cutoff date: {cutoff_date.isoformat()} (retention: {retention_days} days)") + + terminal_statuses = list(TERMINAL_SUCCESS_STATUSES) + if not keep_failed: + terminal_statuses.extend(TERMINAL_FAILED_STATUSES) + + statement = select(WorkflowExecution.id).where( + col(WorkflowExecution.started) < cutoff_date, + col(WorkflowExecution.status).in_(terminal_statuses), + ) + + if tenant_id: + statement = statement.where(WorkflowExecution.tenant_id == tenant_id) + + total_executions_purged = 0 + total_logs_purged = 0 + + while True: + batch_stmt = statement.limit(batch_size) + execution_ids = session.exec(batch_stmt).all() + + if not execution_ids: + break + + logger.info(f"Found {len(execution_ids)} executions to purge in current batch...") + + if dry_run: + total_executions_purged += len(execution_ids) + log_count_stmt = select(WorkflowExecutionLog.id).where( + col(WorkflowExecutionLog.workflow_execution_id).in_(execution_ids) + ) + logs_in_batch = len(session.exec(log_count_stmt).all()) + total_logs_purged += logs_in_batch + break + + # Delete dependent WorkflowExecutionLog entries first + del_logs_stmt = select(WorkflowExecutionLog).where( + col(WorkflowExecutionLog.workflow_execution_id).in_(execution_ids) + ) + logs = session.exec(del_logs_stmt).all() + for log in logs: + session.delete(log) + total_logs_purged += len(logs) + + # Delete WorkflowExecution entries + del_exec_stmt = select(WorkflowExecution).where( + col(WorkflowExecution.id).in_(execution_ids) + ) + execs = session.exec(del_exec_stmt).all() + for ex in execs: + session.delete(ex) + total_executions_purged += len(execs) + + session.commit() + logger.info(f"Purged batch: {len(execs)} executions, {len(logs)} logs.") + + summary = { + "executions_purged": total_executions_purged, + "logs_purged": total_logs_purged, + "dry_run": dry_run, + } + logger.info(f"Cleanup summary: {summary}") + return summary + + +def main(): + parser = argparse.ArgumentParser(description="Purge historical Keep workflow executions and logs.") + parser.add_argument("--retention-days", type=int, default=30, help="Purge executions older than N days.") + parser.add_argument("--batch-size", type=int, default=500, help="Batch size for deletion chunks.") + parser.add_argument("--dry-run", action="store_true", help="Simulate without deleting records.") + parser.add_argument("--tenant-id", type=str, default=None, help="Filter by specific tenant.") + parser.add_argument("--keep-failed", action="store_true", help="Preserve failed/timeout executions.") + args = parser.parse_args() + + with get_session() as session: + purge_workflow_executions( + session=session, + retention_days=args.retention_days, + batch_size=args.batch_size, + dry_run=args.dry_run, + tenant_id=args.tenant_id, + keep_failed=args.keep_failed, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_cleanup_workflow_executions.py b/tests/test_cleanup_workflow_executions.py new file mode 100644 index 0000000000..ff4dd9d53e --- /dev/null +++ b/tests/test_cleanup_workflow_executions.py @@ -0,0 +1,124 @@ +import unittest +from datetime import datetime, timedelta, timezone +from sqlmodel import Session, SQLModel, create_engine +from keep.api.models.db.workflow import WorkflowExecution, WorkflowExecutionLog +from scripts.cleanup_workflow_executions import purge_workflow_executions + + +class TestCleanupWorkflowExecutions(unittest.TestCase): + def setUp(self): + self.engine = create_engine("sqlite:///:memory:") + SQLModel.metadata.create_all(self.engine) + self.session = Session(self.engine) + + old_time = datetime.now(timezone.utc) - timedelta(days=45) + fresh_time = datetime.now(timezone.utc) - timedelta(days=5) + + # 1. Old successful execution + ex1 = WorkflowExecution( + id="exec-old-success", + workflow_id="wf-1", + tenant_id="tenant-a", + status="success", + started=old_time, + execution_number=1, + ) + self.session.add(ex1) + self.session.add(WorkflowExecutionLog( + workflow_execution_id="exec-old-success", + timestamp=old_time, + message="Step completed", + context={"step": 1}, + )) + + # 2. Old error execution + ex2 = WorkflowExecution( + id="exec-old-error", + workflow_id="wf-1", + tenant_id="tenant-a", + status="error", + started=old_time, + execution_number=2, + ) + self.session.add(ex2) + self.session.add(WorkflowExecutionLog( + workflow_execution_id="exec-old-error", + timestamp=old_time, + message="Step failed", + context={"step": 2}, + )) + + # 3. Old running/in-progress execution (must NEVER be purged) + ex3 = WorkflowExecution( + id="exec-old-running", + workflow_id="wf-1", + tenant_id="tenant-a", + status="in_progress", + started=old_time, + execution_number=3, + ) + self.session.add(ex3) + + # 4. Fresh successful execution (< 30 days) + ex4 = WorkflowExecution( + id="exec-fresh-success", + workflow_id="wf-1", + tenant_id="tenant-a", + status="success", + started=fresh_time, + execution_number=4, + ) + self.session.add(ex4) + + self.session.commit() + + def tearDown(self): + self.session.close() + + def test_dry_run_does_not_delete(self): + res = purge_workflow_executions( + session=self.session, + retention_days=30, + dry_run=True, + ) + self.assertEqual(res["executions_purged"], 2) + self.assertEqual(res["logs_purged"], 2) + + rem_execs = self.session.query(WorkflowExecution).count() + self.assertEqual(rem_execs, 4) + + def test_purge_removes_old_terminal_records(self): + res = purge_workflow_executions( + session=self.session, + retention_days=30, + dry_run=False, + ) + self.assertEqual(res["executions_purged"], 2) + self.assertEqual(res["logs_purged"], 2) + + remaining_ids = [e.id for e in self.session.query(WorkflowExecution).all()] + self.assertIn("exec-old-running", remaining_ids) + self.assertIn("exec-fresh-success", remaining_ids) + self.assertNotIn("exec-old-success", remaining_ids) + self.assertNotIn("exec-old-error", remaining_ids) + + remaining_logs = self.session.query(WorkflowExecutionLog).count() + self.assertEqual(remaining_logs, 0) + + def test_keep_failed_preserves_errors(self): + res = purge_workflow_executions( + session=self.session, + retention_days=30, + dry_run=False, + keep_failed=True, + ) + self.assertEqual(res["executions_purged"], 1) + + remaining_ids = [e.id for e in self.session.query(WorkflowExecution).all()] + self.assertIn("exec-old-error", remaining_ids) + self.assertIn("exec-old-running", remaining_ids) + self.assertIn("exec-fresh-success", remaining_ids) + + +if __name__ == "__main__": + unittest.main() From db31adc66b1527fa50689ab75b8e9505e65e29d5 Mon Sep 17 00:00:00 2001 From: Alex Granbell Date: Thu, 12 Nov 2026 11:59:02 +0400 Subject: [PATCH 3/3] fix(retention): address review feedback on session sync, associations, full dry-run counts, and input validation (#6942) --- scripts/cleanup_workflow_executions.py | 507 +++++++++++++++++++--- tests/test_cleanup_workflow_executions.py | 272 ++++++++++-- 2 files changed, 687 insertions(+), 92 deletions(-) diff --git a/scripts/cleanup_workflow_executions.py b/scripts/cleanup_workflow_executions.py index 866a25d3f1..db7579d142 100644 --- a/scripts/cleanup_workflow_executions.py +++ b/scripts/cleanup_workflow_executions.py @@ -2,36 +2,214 @@ """ scripts/cleanup_workflow_executions.py -PostgreSQL (and MySQL/SQLite) supported workflow execution and log retention cleanup. -Addresses #6942 (PostgreSQL TOAST/heap storage exhaustion on WorkflowExecution & WorkflowExecutionLog). - -Supports: -- --retention-days: purge completed executions older than N days (default: 30) -- --batch-size: purge in bounded chunks to avoid lock contention & WAL spikes (default: 500) -- --dry-run: simulate count of executions and logs that would be purged without deleting -- --tenant-id: limit purge to a specific tenant (default: all tenants) -- --keep-failed: preserve failed/error executions for post-mortem debugging +Database-safe workflow execution and log retention cleanup for Keep (#6942). +Addresses PostgreSQL TOAST/heap storage growth on `workflowexecution` (large JSON `results`) +and `workflowexecutionlog`. + +Key safety & operational properties: +- Uses `get_session_sync()` context management for synchronous CLI execution. +- Validates `retention_days`, `failed_retention_days`, `batch_size`, `log_batch_size`, and + `max_runtime_seconds` before any database access. +- Selects only `WorkflowExecution.id` and deletes via bulk SQL `DELETE` statements so large + `WorkflowExecution.results` TOAST payloads and full `WorkflowExecutionLog` objects are + never loaded into Python memory. +- Deletes dependent rows in `WorkflowExecutionLog` (bounded by `log_batch_size`), + `WorkflowToAlertExecution`, and `WorkflowToIncidentExecution` before deleting parent + `WorkflowExecution` rows, preserving `Alert` and `Incident` records. +- In `--dry-run` mode (CLI default unless `--apply` is specified), reports complete aggregate + counts across the entire eligible dataset without truncating at the first batch. +- Includes `providers_not_configured` among terminal failure statuses and never assumes + `WorkflowExecution.is_running` is a boolean (Keep assigns random positive integers on finish). +- Optional `--preserve-latest-per-workflow` retains the latest non-test successful execution + and the scheduler-relevant last completed execution per `(tenant_id, workflow_id)`. + +Operational & PostgreSQL Prerequisites: +1. Index prerequisite: On large installations where `workflowexecutionlog` lacks an index on + `workflow_execution_id`, build it concurrently before running batched deletes: + `CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_workflowexecutionlog_workflow_execution_id ON workflowexecutionlog (workflow_execution_id);` + This script deliberately does not create indexes automatically. +2. Disk-space expectations: Standard `DELETE` + `VACUUM` marks dead tuples/TOAST pages reusable + within PostgreSQL (`pg_freespace`) rather than immediately shrinking OS filesystem usage. + Reclaiming OS disk space requires a separately planned `VACUUM FULL` or `pg_repack` maintenance window. +3. Checkpoint behavior: Keep's `get_previous_execution_id` and `get_last_workflow_execution_by_workflow_id` + bound history lookups to 1 day (`timedelta(days=1)`), while `get_last_completed_execution` inspects the + latest non-test execution (`status IN ('success', 'error', 'providers_not_configured')`). Enabling + `--preserve-latest-per-workflow` retains both the latest non-test success and the latest scheduler-relevant + completed record per workflow regardless of age. """ +from __future__ import annotations + import argparse import logging import os import sys +import time from datetime import datetime, timedelta, timezone +from typing import Optional, Sequence -from sqlmodel import Session, select, col +from sqlalchemy import and_, delete, func, or_ +from sqlmodel import Session, col, select # Adjust path for keep imports sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) -from keep.api.core.db import get_session -from keep.api.models.db.workflow import WorkflowExecution, WorkflowExecutionLog +from keep.api.core.db import get_session_sync +from keep.api.models.db.workflow import ( + WorkflowExecution, + WorkflowExecutionLog, + WorkflowToAlertExecution, + WorkflowToIncidentExecution, +) logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) -TERMINAL_SUCCESS_STATUSES = ["success", "passed"] -TERMINAL_FAILED_STATUSES = ["error", "timeout", "failed", "cancelled"] +TERMINAL_SUCCESS_STATUSES = ("success", "passed") +TERMINAL_FAILED_STATUSES = ( + "error", + "timeout", + "failed", + "cancelled", + "providers_not_configured", +) +SCHEDULER_COMPLETED_STATUSES = ("success", "error", "providers_not_configured") + + +def _validate_parameters( + retention_days: int, + batch_size: int, + failed_retention_days: int | None = None, + log_batch_size: int = 5000, + max_runtime_seconds: float | None = None, +) -> None: + if not isinstance(retention_days, int) or isinstance(retention_days, bool) or retention_days < 1: + raise ValueError(f"retention_days must be an integer >= 1 (got {retention_days!r})") + if failed_retention_days is not None and ( + not isinstance(failed_retention_days, int) + or isinstance(failed_retention_days, bool) + or failed_retention_days < 1 + ): + raise ValueError( + f"failed_retention_days must be an integer >= 1 when specified (got {failed_retention_days!r})" + ) + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1: + raise ValueError(f"batch_size must be an integer >= 1 (got {batch_size!r})") + if not isinstance(log_batch_size, int) or isinstance(log_batch_size, bool) or log_batch_size < 1: + raise ValueError(f"log_batch_size must be an integer >= 1 (got {log_batch_size!r})") + if max_runtime_seconds is not None and max_runtime_seconds <= 0: + raise ValueError(f"max_runtime_seconds must be > 0 when specified (got {max_runtime_seconds!r})") + + +def get_protected_execution_ids( + session: Session, + tenant_id: str | None = None, +) -> set[str]: + """ + Identify execution IDs that should be preserved per (tenant_id, workflow_id): + 1. Latest non-test successful execution (`status IN TERMINAL_SUCCESS_STATUSES`). + 2. Scheduler-relevant last completed non-test execution (`status IN SCHEDULER_COMPLETED_STATUSES`, + ordered by `execution_number DESC, started DESC`). + """ + protected_ids: set[str] = set() + + base_filters = [col(WorkflowExecution.is_test_run) == False] # noqa: E712 + if tenant_id: + base_filters.append(WorkflowExecution.tenant_id == tenant_id) + + workflows_stmt = ( + select(WorkflowExecution.tenant_id, WorkflowExecution.workflow_id) + .where(*base_filters) + .distinct() + ) + tenant_workflows = session.exec(workflows_stmt).all() + + for t_id, wf_id in tenant_workflows: + # 1. Latest non-test success + latest_success_id = session.exec( + select(WorkflowExecution.id) + .where( + WorkflowExecution.tenant_id == t_id, + WorkflowExecution.workflow_id == wf_id, + col(WorkflowExecution.is_test_run) == False, # noqa: E712 + col(WorkflowExecution.status).in_(TERMINAL_SUCCESS_STATUSES), + ) + .order_by( + col(WorkflowExecution.started).desc(), + col(WorkflowExecution.execution_number).desc(), + ) + .limit(1) + ).first() + if latest_success_id: + protected_ids.add(latest_success_id) + + # 2. Latest scheduler-relevant completed record + latest_completed_id = session.exec( + select(WorkflowExecution.id) + .where( + WorkflowExecution.tenant_id == t_id, + WorkflowExecution.workflow_id == wf_id, + col(WorkflowExecution.is_test_run) == False, # noqa: E712 + col(WorkflowExecution.status).in_(SCHEDULER_COMPLETED_STATUSES), + ) + .order_by( + col(WorkflowExecution.execution_number).desc(), + col(WorkflowExecution.started).desc(), + ) + .limit(1) + ).first() + if latest_completed_id: + protected_ids.add(latest_completed_id) + + return protected_ids + + +def _build_eligible_ids_statement( + now: datetime, + retention_days: int, + failed_retention_days: int | None, + tenant_id: str | None, + keep_failed: bool, + protected_ids: set[str], +): + success_cutoff = now - timedelta(days=retention_days) + effective_failed_days = ( + failed_retention_days if failed_retention_days is not None else retention_days + ) + failed_cutoff = now - timedelta(days=effective_failed_days) + + if keep_failed: + status_time_condition = and_( + col(WorkflowExecution.status).in_(TERMINAL_SUCCESS_STATUSES), + col(WorkflowExecution.started) < success_cutoff, + ) + elif failed_retention_days is not None and failed_retention_days != retention_days: + status_time_condition = or_( + and_( + col(WorkflowExecution.status).in_(TERMINAL_SUCCESS_STATUSES), + col(WorkflowExecution.started) < success_cutoff, + ), + and_( + col(WorkflowExecution.status).in_(TERMINAL_FAILED_STATUSES), + col(WorkflowExecution.started) < failed_cutoff, + ), + ) + else: + all_terminal = list(TERMINAL_SUCCESS_STATUSES) + list(TERMINAL_FAILED_STATUSES) + status_time_condition = and_( + col(WorkflowExecution.status).in_(all_terminal), + col(WorkflowExecution.started) < success_cutoff, + ) + + stmt = select(WorkflowExecution.id).where(status_time_condition) + + if tenant_id: + stmt = stmt.where(WorkflowExecution.tenant_id == tenant_id) + + if protected_ids: + stmt = stmt.where(col(WorkflowExecution.id).notin_(protected_ids)) + + return stmt.order_by(col(WorkflowExecution.started).asc(), col(WorkflowExecution.id).asc()) def purge_workflow_executions( @@ -41,90 +219,287 @@ def purge_workflow_executions( dry_run: bool = False, tenant_id: str | None = None, keep_failed: bool = False, -) -> dict[str, int]: - cutoff_date = datetime.now(timezone.utc) - timedelta(days=retention_days) - logger.info(f"Purge cutoff date: {cutoff_date.isoformat()} (retention: {retention_days} days)") + failed_retention_days: int | None = None, + log_batch_size: int = 5000, + preserve_latest_per_workflow: bool = False, + max_runtime_seconds: float | None = None, +) -> dict[str, int | bool]: + _validate_parameters( + retention_days=retention_days, + batch_size=batch_size, + failed_retention_days=failed_retention_days, + log_batch_size=log_batch_size, + max_runtime_seconds=max_runtime_seconds, + ) - terminal_statuses = list(TERMINAL_SUCCESS_STATUSES) - if not keep_failed: - terminal_statuses.extend(TERMINAL_FAILED_STATUSES) + now = datetime.now(timezone.utc) + start_monotonic = time.monotonic() - statement = select(WorkflowExecution.id).where( - col(WorkflowExecution.started) < cutoff_date, - col(WorkflowExecution.status).in_(terminal_statuses), + protected_ids: set[str] = set() + if preserve_latest_per_workflow: + protected_ids = get_protected_execution_ids(session=session, tenant_id=tenant_id) + + statement = _build_eligible_ids_statement( + now=now, + retention_days=retention_days, + failed_retention_days=failed_retention_days, + tenant_id=tenant_id, + keep_failed=keep_failed, + protected_ids=protected_ids, ) - if tenant_id: - statement = statement.where(WorkflowExecution.tenant_id == tenant_id) + if dry_run: + eligible_subq = statement.subquery() + eligible_id_select = select(eligible_subq.c.id) + + total_executions = int( + session.exec(select(func.count()).select_from(eligible_subq)).one() or 0 + ) + total_logs = int( + session.exec( + select(func.count(WorkflowExecutionLog.id)).where( + col(WorkflowExecutionLog.workflow_execution_id).in_(eligible_id_select) + ) + ).one() + or 0 + ) + total_alert_assocs = int( + session.exec( + select(func.count(WorkflowToAlertExecution.id)).where( + col(WorkflowToAlertExecution.workflow_execution_id).in_(eligible_id_select) + ) + ).one() + or 0 + ) + total_incident_assocs = int( + session.exec( + select(func.count(WorkflowToIncidentExecution.id)).where( + col(WorkflowToIncidentExecution.workflow_execution_id).in_( + eligible_id_select + ) + ) + ).one() + or 0 + ) + + summary: dict[str, int | bool] = { + "executions_purged": total_executions, + "logs_purged": total_logs, + "alert_associations_purged": total_alert_assocs, + "incident_associations_purged": total_incident_assocs, + "batches_completed": 0, + "dry_run": True, + "time_limit_reached": False, + } + logger.info(f"Dry-run summary (full eligible scope): {summary}") + return summary total_executions_purged = 0 total_logs_purged = 0 + total_alert_assocs_purged = 0 + total_incident_assocs_purged = 0 + batches_completed = 0 + time_limit_reached = False while True: - batch_stmt = statement.limit(batch_size) - execution_ids = session.exec(batch_stmt).all() + if ( + max_runtime_seconds is not None + and (time.monotonic() - start_monotonic) >= max_runtime_seconds + ): + time_limit_reached = True + logger.info("Reached max_runtime_seconds limit; stopping cleanly after committed batches.") + break + execution_ids = list(session.exec(statement.limit(batch_size)).all()) if not execution_ids: break logger.info(f"Found {len(execution_ids)} executions to purge in current batch...") - if dry_run: - total_executions_purged += len(execution_ids) - log_count_stmt = select(WorkflowExecutionLog.id).where( - col(WorkflowExecutionLog.workflow_execution_id).in_(execution_ids) + try: + # 1. Delete dependent WorkflowExecutionLog entries in bounded chunks + batch_logs_deleted = 0 + while True: + log_ids = list( + session.exec( + select(WorkflowExecutionLog.id) + .where( + col(WorkflowExecutionLog.workflow_execution_id).in_(execution_ids) + ) + .limit(log_batch_size) + ).all() + ) + if not log_ids: + break + session.exec( + delete(WorkflowExecutionLog).where( + col(WorkflowExecutionLog.id).in_(log_ids) + ) + ) + batch_logs_deleted += len(log_ids) + + # 2. Delete WorkflowToAlertExecution associations (preserving Alert rows) + alert_res = session.exec( + delete(WorkflowToAlertExecution).where( + col(WorkflowToAlertExecution.workflow_execution_id).in_(execution_ids) + ) ) - logs_in_batch = len(session.exec(log_count_stmt).all()) - total_logs_purged += logs_in_batch - break + batch_alert_assocs = int(getattr(alert_res, "rowcount", 0) or 0) - # Delete dependent WorkflowExecutionLog entries first - del_logs_stmt = select(WorkflowExecutionLog).where( - col(WorkflowExecutionLog.workflow_execution_id).in_(execution_ids) - ) - logs = session.exec(del_logs_stmt).all() - for log in logs: - session.delete(log) - total_logs_purged += len(logs) - - # Delete WorkflowExecution entries - del_exec_stmt = select(WorkflowExecution).where( - col(WorkflowExecution.id).in_(execution_ids) - ) - execs = session.exec(del_exec_stmt).all() - for ex in execs: - session.delete(ex) - total_executions_purged += len(execs) + # 3. Delete WorkflowToIncidentExecution associations (preserving Incident rows) + incident_res = session.exec( + delete(WorkflowToIncidentExecution).where( + col(WorkflowToIncidentExecution.workflow_execution_id).in_(execution_ids) + ) + ) + batch_incident_assocs = int(getattr(incident_res, "rowcount", 0) or 0) - session.commit() - logger.info(f"Purged batch: {len(execs)} executions, {len(logs)} logs.") + # 4. Delete parent WorkflowExecution rows by ID without loading `results` JSON + exec_res = session.exec( + delete(WorkflowExecution).where( + col(WorkflowExecution.id).in_(execution_ids) + ) + ) + batch_execs_deleted = int( + getattr(exec_res, "rowcount", len(execution_ids)) or len(execution_ids) + ) + + session.commit() + except Exception: + session.rollback() + logger.exception("Error purging batch; transaction rolled back.") + raise + + total_executions_purged += batch_execs_deleted + total_logs_purged += batch_logs_deleted + total_alert_assocs_purged += batch_alert_assocs + total_incident_assocs_purged += batch_incident_assocs + batches_completed += 1 + + logger.info( + f"Purged batch #{batches_completed}: {batch_execs_deleted} executions, " + f"{batch_logs_deleted} logs, {batch_alert_assocs} alert links, " + f"{batch_incident_assocs} incident links." + ) summary = { "executions_purged": total_executions_purged, "logs_purged": total_logs_purged, - "dry_run": dry_run, + "alert_associations_purged": total_alert_assocs_purged, + "incident_associations_purged": total_incident_assocs_purged, + "batches_completed": batches_completed, + "dry_run": False, + "time_limit_reached": time_limit_reached, } logger.info(f"Cleanup summary: {summary}") return summary -def main(): - parser = argparse.ArgumentParser(description="Purge historical Keep workflow executions and logs.") - parser.add_argument("--retention-days", type=int, default=30, help="Purge executions older than N days.") - parser.add_argument("--batch-size", type=int, default=500, help="Batch size for deletion chunks.") - parser.add_argument("--dry-run", action="store_true", help="Simulate without deleting records.") - parser.add_argument("--tenant-id", type=str, default=None, help="Filter by specific tenant.") - parser.add_argument("--keep-failed", action="store_true", help="Preserve failed/timeout executions.") - args = parser.parse_args() +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Purge historical Keep workflow executions and logs safely." + ) + parser.add_argument( + "--retention-days", + type=int, + default=30, + help="Purge completed executions older than N days (default: 30, must be >= 1).", + ) + parser.add_argument( + "--failed-retention-days", + type=int, + default=None, + help="Optional separate retention window (in days, >= 1) for failed/error executions.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=500, + help="Number of workflow executions to purge per transaction batch (default: 500).", + ) + parser.add_argument( + "--log-batch-size", + type=int, + default=5000, + help="Maximum WorkflowExecutionLog rows deleted per statement chunk (default: 5000).", + ) + parser.add_argument( + "--max-runtime-seconds", + type=float, + default=None, + help="Optional maximum runtime in seconds before stopping cleanly between batches.", + ) + + mode_group = parser.add_mutually_exclusive_group() + mode_group.add_argument( + "--dry-run", + action="store_true", + help="Report full eligible counts without deleting records (default mode unless --apply is set).", + ) + mode_group.add_argument( + "--apply", + action="store_true", + help="Execute actual deletion batches in the database.", + ) + + tenant_group = parser.add_mutually_exclusive_group(required=True) + tenant_group.add_argument( + "--tenant-id", + type=str, + default=None, + help="Filter cleanup to a specific tenant ID.", + ) + tenant_group.add_argument( + "--all-tenants", + action="store_true", + help="Explicitly run cleanup across all tenants.", + ) + + parser.add_argument( + "--keep-failed", + action="store_true", + help="Preserve failed/timeout/providers_not_configured executions.", + ) + parser.add_argument( + "--preserve-latest-per-workflow", + action="store_true", + help="Preserve the latest non-test success and scheduler-relevant completed execution per workflow.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> dict[str, int | bool]: + parser = build_arg_parser() + args = parser.parse_args(argv) + + # Validate parameters before opening a database session + _validate_parameters( + retention_days=args.retention_days, + batch_size=args.batch_size, + failed_retention_days=args.failed_retention_days, + log_batch_size=args.log_batch_size, + max_runtime_seconds=args.max_runtime_seconds, + ) + + # Default to dry-run unless --apply is explicitly provided + effective_dry_run = True if not args.apply else False + if args.dry_run: + effective_dry_run = True + + effective_tenant_id = None if args.all_tenants else args.tenant_id - with get_session() as session: - purge_workflow_executions( + with get_session_sync() as session: + return purge_workflow_executions( session=session, retention_days=args.retention_days, batch_size=args.batch_size, - dry_run=args.dry_run, - tenant_id=args.tenant_id, + dry_run=effective_dry_run, + tenant_id=effective_tenant_id, keep_failed=args.keep_failed, + failed_retention_days=args.failed_retention_days, + log_batch_size=args.log_batch_size, + preserve_latest_per_workflow=args.preserve_latest_per_workflow, + max_runtime_seconds=args.max_runtime_seconds, ) diff --git a/tests/test_cleanup_workflow_executions.py b/tests/test_cleanup_workflow_executions.py index ff4dd9d53e..73bbaccb92 100644 --- a/tests/test_cleanup_workflow_executions.py +++ b/tests/test_cleanup_workflow_executions.py @@ -1,61 +1,132 @@ import unittest from datetime import datetime, timedelta, timezone -from sqlmodel import Session, SQLModel, create_engine -from keep.api.models.db.workflow import WorkflowExecution, WorkflowExecutionLog -from scripts.cleanup_workflow_executions import purge_workflow_executions +from unittest.mock import patch + +from sqlalchemy import event +from sqlmodel import Session, SQLModel, col, create_engine, select + +from keep.api.models.db.tenant import Tenant +from keep.api.models.db.workflow import ( + Workflow, + WorkflowExecution, + WorkflowExecutionLog, + WorkflowToAlertExecution, + WorkflowToIncidentExecution, +) +from scripts.cleanup_workflow_executions import ( + get_protected_execution_ids, + main, + purge_workflow_executions, +) class TestCleanupWorkflowExecutions(unittest.TestCase): def setUp(self): self.engine = create_engine("sqlite:///:memory:") + + @event.listens_for(self.engine, "connect") + def _enable_sqlite_fks(dbapi_connection, connection_record): + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + SQLModel.metadata.create_all(self.engine) self.session = Session(self.engine) - old_time = datetime.now(timezone.utc) - timedelta(days=45) - fresh_time = datetime.now(timezone.utc) - timedelta(days=5) + self.session.add(Tenant(id="tenant-a", name="Tenant A")) + self.session.add( + Workflow( + id="wf-1", + tenant_id="tenant-a", + name="wf-1", + description="test workflow", + created_by="unit-test", + workflow_raw="{}", + ) + ) + self.session.add( + Workflow( + id="wf-dormant", + tenant_id="tenant-a", + name="wf-dormant", + description="dormant workflow", + created_by="unit-test", + workflow_raw="{}", + ) + ) + self.session.commit() + + self.old_time = datetime.now(timezone.utc) - timedelta(days=45) + self.fresh_time = datetime.now(timezone.utc) - timedelta(days=5) - # 1. Old successful execution + # 1. Old successful execution with logs + alert & incident associations ex1 = WorkflowExecution( id="exec-old-success", workflow_id="wf-1", tenant_id="tenant-a", + triggered_by="scheduler", status="success", - started=old_time, + started=self.old_time, execution_number=1, + is_running=42123, # Keep sets random positive int on finish + is_test_run=False, ) self.session.add(ex1) - self.session.add(WorkflowExecutionLog( - workflow_execution_id="exec-old-success", - timestamp=old_time, - message="Step completed", - context={"step": 1}, - )) + self.session.add( + WorkflowExecutionLog( + workflow_execution_id="exec-old-success", + timestamp=self.old_time, + message="Step completed", + context={"step": 1}, + ) + ) + self.session.add( + WorkflowToAlertExecution( + workflow_execution_id="exec-old-success", + alert_fingerprint="fp-alert-1", + event_id="evt-1", + ) + ) + self.session.add( + WorkflowToIncidentExecution( + workflow_execution_id="exec-old-success", + incident_id="inc-1", + ) + ) # 2. Old error execution ex2 = WorkflowExecution( id="exec-old-error", workflow_id="wf-1", tenant_id="tenant-a", + triggered_by="alert", status="error", - started=old_time, + started=self.old_time, execution_number=2, + is_running=98765, + is_test_run=False, ) self.session.add(ex2) - self.session.add(WorkflowExecutionLog( - workflow_execution_id="exec-old-error", - timestamp=old_time, - message="Step failed", - context={"step": 2}, - )) + self.session.add( + WorkflowExecutionLog( + workflow_execution_id="exec-old-error", + timestamp=self.old_time, + message="Step failed", + context={"step": 2}, + ) + ) # 3. Old running/in-progress execution (must NEVER be purged) ex3 = WorkflowExecution( id="exec-old-running", workflow_id="wf-1", tenant_id="tenant-a", + triggered_by="manual", status="in_progress", - started=old_time, + started=self.old_time, execution_number=3, + is_running=1, + is_test_run=False, ) self.session.add(ex3) @@ -64,9 +135,12 @@ def setUp(self): id="exec-fresh-success", workflow_id="wf-1", tenant_id="tenant-a", + triggered_by="scheduler", status="success", - started=fresh_time, + started=self.fresh_time, execution_number=4, + is_running=55111, + is_test_run=False, ) self.session.add(ex4) @@ -75,7 +149,7 @@ def setUp(self): def tearDown(self): self.session.close() - def test_dry_run_does_not_delete(self): + def test_dry_run_does_not_delete_and_reports_associations(self): res = purge_workflow_executions( session=self.session, retention_days=30, @@ -83,11 +157,62 @@ def test_dry_run_does_not_delete(self): ) self.assertEqual(res["executions_purged"], 2) self.assertEqual(res["logs_purged"], 2) + self.assertEqual(res["alert_associations_purged"], 1) + self.assertEqual(res["incident_associations_purged"], 1) rem_execs = self.session.query(WorkflowExecution).count() self.assertEqual(rem_execs, 4) - def test_purge_removes_old_terminal_records(self): + def test_dry_run_reports_full_scope_across_multiple_batches(self): + # Add 3 more old terminal executions -> 5 eligible total + for i in range(3): + eid = f"exec-extra-{i}" + self.session.add( + WorkflowExecution( + id=eid, + workflow_id="wf-1", + tenant_id="tenant-a", + triggered_by="scheduler", + status="providers_not_configured" if i == 0 else "success", + started=self.old_time - timedelta(days=i + 1), + execution_number=10 + i, + is_running=100 + i, + is_test_run=False, + ) + ) + self.session.add( + WorkflowExecutionLog( + workflow_execution_id=eid, + timestamp=self.old_time, + message=f"Extra log {i}", + context={}, + ) + ) + self.session.commit() + + # With batch_size=2, dry-run must report all 5 eligible executions and 5 logs + res_dry = purge_workflow_executions( + session=self.session, + retention_days=30, + batch_size=2, + dry_run=True, + ) + self.assertEqual(res_dry["executions_purged"], 5) + self.assertEqual(res_dry["logs_purged"], 5) + + # Apply mode with batch_size=2 must iterate across 3 batches to purge all 5 + res_apply = purge_workflow_executions( + session=self.session, + retention_days=30, + batch_size=2, + log_batch_size=1, + dry_run=False, + ) + self.assertEqual(res_apply["executions_purged"], 5) + self.assertEqual(res_apply["logs_purged"], 5) + self.assertEqual(res_apply["batches_completed"], 3) + + def test_purge_removes_old_terminal_records_and_associations(self): res = purge_workflow_executions( session=self.session, retention_days=30, @@ -95,6 +220,8 @@ def test_purge_removes_old_terminal_records(self): ) self.assertEqual(res["executions_purged"], 2) self.assertEqual(res["logs_purged"], 2) + self.assertEqual(res["alert_associations_purged"], 1) + self.assertEqual(res["incident_associations_purged"], 1) remaining_ids = [e.id for e in self.session.query(WorkflowExecution).all()] self.assertIn("exec-old-running", remaining_ids) @@ -102,8 +229,9 @@ def test_purge_removes_old_terminal_records(self): self.assertNotIn("exec-old-success", remaining_ids) self.assertNotIn("exec-old-error", remaining_ids) - remaining_logs = self.session.query(WorkflowExecutionLog).count() - self.assertEqual(remaining_logs, 0) + self.assertEqual(self.session.query(WorkflowExecutionLog).count(), 0) + self.assertEqual(self.session.query(WorkflowToAlertExecution).count(), 0) + self.assertEqual(self.session.query(WorkflowToIncidentExecution).count(), 0) def test_keep_failed_preserves_errors(self): res = purge_workflow_executions( @@ -119,6 +247,98 @@ def test_keep_failed_preserves_errors(self): self.assertIn("exec-old-running", remaining_ids) self.assertIn("exec-fresh-success", remaining_ids) + def test_invalid_retention_and_batch_size_rejected_before_db_access(self): + for bad_retention in (-1, 0, False): + with self.assertRaises(ValueError): + purge_workflow_executions( + session=self.session, + retention_days=bad_retention, # type: ignore[arg-type] + dry_run=False, + ) + for bad_batch in (-5, 0): + with self.assertRaises(ValueError): + purge_workflow_executions( + session=self.session, + retention_days=30, + batch_size=bad_batch, + dry_run=False, + ) + # Ensure fresh execution was untouched + self.assertEqual(self.session.query(WorkflowExecution).count(), 4) + + def test_preserve_latest_per_workflow_protects_dormant_workflow_checkpoint(self): + # Dormant workflow wf-dormant has only old executions (#1 at 60d, #2 at 50d) + self.session.add( + WorkflowExecution( + id="dormant-1", + workflow_id="wf-dormant", + tenant_id="tenant-a", + triggered_by="scheduler", + status="success", + started=datetime.now(timezone.utc) - timedelta(days=60), + execution_number=1, + is_running=111, + is_test_run=False, + ) + ) + self.session.add( + WorkflowExecution( + id="dormant-2-latest", + workflow_id="wf-dormant", + tenant_id="tenant-a", + triggered_by="scheduler", + status="success", + started=datetime.now(timezone.utc) - timedelta(days=50), + execution_number=2, + is_running=222, + is_test_run=False, + ) + ) + self.session.commit() + + protected = get_protected_execution_ids(self.session, tenant_id="tenant-a") + self.assertIn("dormant-2-latest", protected) + self.assertNotIn("dormant-1", protected) + + res = purge_workflow_executions( + session=self.session, + retention_days=30, + dry_run=False, + preserve_latest_per_workflow=True, + ) + # Purges exec-old-success, exec-old-error, and dormant-1 (3 total), keeping dormant-2-latest + self.assertEqual(res["executions_purged"], 3) + remaining_ids = [e.id for e in self.session.query(WorkflowExecution).all()] + self.assertIn("dormant-2-latest", remaining_ids) + self.assertNotIn("dormant-1", remaining_ids) + + def test_cli_main_uses_get_session_sync_and_defaults_to_dry_run(self): + class _DummySessionCtx: + def __init__(self, sess): + self._sess = sess + + def __enter__(self): + return self._sess + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + with patch( + "scripts.cleanup_workflow_executions.get_session_sync", + return_value=_DummySessionCtx(self.session), + ): + # Default CLI mode (without --apply) must be dry_run=True + res_dry = main(["--retention-days", "30", "--tenant-id", "tenant-a"]) + self.assertTrue(res_dry["dry_run"]) + self.assertEqual(res_dry["executions_purged"], 2) + self.assertEqual(self.session.query(WorkflowExecution).count(), 4) + + # Explicit --apply mode deletes the 2 eligible records + res_apply = main(["--retention-days", "30", "--all-tenants", "--apply"]) + self.assertFalse(res_apply["dry_run"]) + self.assertEqual(res_apply["executions_purged"], 2) + self.assertEqual(self.session.query(WorkflowExecution).count(), 2) + if __name__ == "__main__": unittest.main()