diff --git a/src/executors/postgres_retention_executor.py b/src/executors/postgres_retention_executor.py new file mode 100644 index 0000000..2ede15e --- /dev/null +++ b/src/executors/postgres_retention_executor.py @@ -0,0 +1,109 @@ +""" +Postgres Retention Executor + +Deletes rows older than a retention window from a table on the shared Postgres +server. Written for sysmon's `check_history`, which grows with every monitoring +check and had no retention at all, but the executor is table-agnostic. + +Connects with the Scheduler's own Postgres credentials and only overrides the +database name. That keeps a second set of credentials out of the stack; the +target database grants `scheduler_user` exactly SELECT and DELETE on the table, +so a bug here can drop old rows but cannot corrupt or forge history. + +Config schema: + { + "database": "sysmon", # defaults to the Scheduler's own database + "table": "check_history", # required + "timestamp_column": "ts", # required + "retention_days": 30, # required, must be >= 1 + "dry_run": false # count what would go, delete nothing + } + +Table and column names cannot be passed as query parameters, so both are +validated against a strict identifier pattern before being interpolated. + +Autovacuum reclaims the space afterwards; this deliberately does not VACUUM, +which would need table ownership the Scheduler intentionally does not have. +""" +import asyncio +import logging +import re +from typing import Any, Dict + +import psycopg2 + +from src.config import Settings + +logger = logging.getLogger(__name__) + +# Deliberately strict: unquoted lowercase identifiers only. Anything needing +# quoting is out of scope and would be a hole in the interpolation below. +IDENTIFIER_RE = re.compile(r"^[a-z_][a-z0-9_]*$") + +MAX_RETENTION_DAYS = 3650 + + +def _validate_identifier(value: str, label: str) -> str: + if not isinstance(value, str) or not IDENTIFIER_RE.match(value): + raise ValueError( + f"invalid {label}: {value!r} (expected an unquoted lowercase identifier)" + ) + return value + + +def _prune(config: Dict[str, Any], settings: Settings) -> str: + table = _validate_identifier(config.get("table", ""), "table") + column = _validate_identifier(config.get("timestamp_column", ""), "timestamp_column") + database = config.get("database") or settings.postgres_db + _validate_identifier(database, "database") + + retention_days = config.get("retention_days") + if not isinstance(retention_days, int) or isinstance(retention_days, bool): + raise ValueError(f"retention_days must be an integer, got {retention_days!r}") + # A zero or negative window would delete everything, including the row the + # check just wrote. Refuse rather than quietly wipe the table. + if retention_days < 1 or retention_days > MAX_RETENTION_DAYS: + raise ValueError( + f"retention_days must be between 1 and {MAX_RETENTION_DAYS}, got {retention_days}" + ) + + dry_run = bool(config.get("dry_run", False)) + cutoff_sql = f"{column} < now() - make_interval(days => %s)" + + conn = psycopg2.connect( + host=settings.postgres_host, + port=settings.postgres_port, + database=database, + user=settings.postgres_user, + password=settings.postgres_password, + connect_timeout=10, + ) + try: + with conn: + with conn.cursor() as cur: + cur.execute(f"SELECT count(*) FROM {table} WHERE {cutoff_sql}", (retention_days,)) + stale = cur.fetchone()[0] + + if dry_run: + logger.info("dry run: %s rows in %s.%s exceed %sd", stale, database, table, retention_days) + return f"dry run: {stale} rows older than {retention_days}d in {database}.{table}" + + if stale == 0: + return f"nothing to prune in {database}.{table} (retention {retention_days}d)" + + cur.execute(f"DELETE FROM {table} WHERE {cutoff_sql}", (retention_days,)) + deleted = cur.rowcount + + cur.execute(f"SELECT count(*) FROM {table}") + remaining = cur.fetchone()[0] + finally: + conn.close() + + logger.info("pruned %s rows from %s.%s, %s remain", deleted, database, table, remaining) + return f"pruned {deleted} rows older than {retention_days}d from {database}.{table}, {remaining} remain" + + +async def execute(config: dict, settings: Settings) -> str: + """Delete rows past the retention window. Returns a one-line summary.""" + # psycopg2 is synchronous; keep it off the scheduler's event loop. + return await asyncio.to_thread(_prune, config, settings) diff --git a/tests/test_postgres_retention_executor.py b/tests/test_postgres_retention_executor.py new file mode 100644 index 0000000..18539a3 --- /dev/null +++ b/tests/test_postgres_retention_executor.py @@ -0,0 +1,151 @@ +""" +Tests for the postgres retention executor. + +Focus is on the guards. The executor interpolates a table and column name +straight into SQL (they cannot be bound as parameters), and it issues DELETEs +against a live table, so the validation in front of both is what keeps a +malformed config from becoming data loss. +""" +from unittest.mock import MagicMock, patch + +import pytest + +from src.config import Settings +from src.executors import postgres_retention_executor as retention + + +@pytest.mark.executor +@pytest.mark.unit +class TestIdentifierValidation: + """Table/column/database names are interpolated, so they must be rejected early.""" + + @pytest.mark.parametrize( + "bad", + [ + "check_history; DROP TABLE users", + 'check_history"', + "check history", + "Check_History", # uppercase would need quoting to resolve + "1_history", + "", + "--comment", + ], + ) + def test_rejects_unsafe_identifiers(self, bad): + with pytest.raises(ValueError): + retention._validate_identifier(bad, "table") + + @pytest.mark.parametrize("good", ["check_history", "ts", "_private", "a1"]) + def test_accepts_plain_identifiers(self, good): + assert retention._validate_identifier(good, "table") == good + + +@pytest.mark.executor +@pytest.mark.unit +class TestRetentionGuards: + """A bad retention window must never reach the database.""" + + def _config(self, **overrides): + config = { + "database": "sysmon", + "table": "check_history", + "timestamp_column": "ts", + "retention_days": 30, + } + config.update(overrides) + return config + + @pytest.mark.parametrize("days", [0, -1, -30, 3651]) + def test_rejects_out_of_range_retention(self, days, test_settings: Settings): + # 0 or negative would delete every row including the one just written. + with patch("psycopg2.connect") as connect: + with pytest.raises(ValueError): + retention._prune(self._config(retention_days=days), test_settings) + connect.assert_not_called() + + @pytest.mark.parametrize("days", ["30", None, 1.5, True]) + def test_rejects_non_integer_retention(self, days, test_settings: Settings): + with patch("psycopg2.connect") as connect: + with pytest.raises(ValueError): + retention._prune(self._config(retention_days=days), test_settings) + connect.assert_not_called() + + def test_rejects_injection_in_table_before_connecting(self, test_settings: Settings): + with patch("psycopg2.connect") as connect: + with pytest.raises(ValueError): + retention._prune( + self._config(table="check_history; DELETE FROM check_history --"), + test_settings, + ) + connect.assert_not_called() + + +@pytest.mark.executor +@pytest.mark.unit +class TestRetentionBehaviour: + """Behaviour against a mocked cursor.""" + + def _mock_conn(self, counts): + cursor = MagicMock() + cursor.fetchone.side_effect = [(c,) for c in counts] + cursor.rowcount = counts[0] if counts else 0 + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value = cursor + conn.__enter__.return_value = conn + return conn, cursor + + def _config(self, **overrides): + config = { + "database": "sysmon", + "table": "check_history", + "timestamp_column": "ts", + "retention_days": 30, + } + config.update(overrides) + return config + + def test_dry_run_does_not_delete(self, test_settings: Settings): + conn, cursor = self._mock_conn([7]) + with patch("psycopg2.connect", return_value=conn): + result = retention._prune(self._config(dry_run=True), test_settings) + + assert "dry run" in result + assert "7" in result + executed = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "DELETE" not in executed.upper() + + def test_no_stale_rows_skips_delete(self, test_settings: Settings): + conn, cursor = self._mock_conn([0]) + with patch("psycopg2.connect", return_value=conn): + result = retention._prune(self._config(), test_settings) + + assert "nothing to prune" in result + executed = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "DELETE" not in executed.upper() + + def test_deletes_and_reports(self, test_settings: Settings): + # count(stale) -> 5, then count(remaining) -> 42 + conn, cursor = self._mock_conn([5, 42]) + cursor.rowcount = 5 + with patch("psycopg2.connect", return_value=conn): + result = retention._prune(self._config(), test_settings) + + assert "pruned 5 rows" in result + assert "42 remain" in result + executed = " ".join(str(c) for c in cursor.execute.call_args_list) + assert "DELETE" in executed.upper() + + def test_defaults_to_scheduler_database(self, test_settings: Settings): + conn, _ = self._mock_conn([0]) + config = self._config() + del config["database"] + with patch("psycopg2.connect", return_value=conn) as connect: + retention._prune(config, test_settings) + assert connect.call_args.kwargs["database"] == test_settings.postgres_db + + @pytest.mark.asyncio + async def test_execute_wraps_prune(self, test_settings: Settings): + conn, _ = self._mock_conn([0]) + with patch("psycopg2.connect", return_value=conn): + result = await retention.execute(self._config(), test_settings) + assert "nothing to prune" in result