Deletes rows past 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 despite the docs promising a 30-day rolling window. Connects with the Scheduler's own credentials and overrides only the database name, so no second set of secrets enters the stack. The target database grants scheduler_user just SELECT and DELETE on the table, so a bug here can drop old rows but cannot corrupt or forge history. Table and column names cannot be bound as query parameters, so both are validated against a strict identifier pattern before interpolation, and a retention window below 1 day is refused rather than silently emptying the table. Co-Authored-By: Claude <noreply@anthropic.com>
152 lines
5.6 KiB
Python
152 lines
5.6 KiB
Python
"""
|
|
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
|