from __future__ import annotations import os import unittest from datetime import date from unittest.mock import patch from arr_processing.contracts import ProcessingRequest from arr_processing.errors import ProcessingError from arr_processing.postgres import ( PostgresProcessingState, ProcessingDatabaseConfig, ) from arr_processing.runner import processing_dispatch_idempotency_key class ScriptedConnection: def __init__(self, rows_by_marker: dict[str, tuple[object, ...] | None]) -> None: self.rows_by_marker = rows_by_marker self.last_query = "" self.executed: list[tuple[str, tuple[object, ...] | None]] = [] self.closed = False def transaction(self): return self def cursor(self): return self def __enter__(self): return self def __exit__(self, _kind, _value, _traceback): return False def execute(self, query: str, params: tuple[object, ...] | None = None) -> None: self.last_query = " ".join(query.split()) self.executed.append((self.last_query, params)) def fetchone(self): if "current_database()" in self.last_query: return ("booking_test",) for marker, row in self.rows_by_marker.items(): if marker in self.last_query: return row return None def close(self) -> None: self.closed = True def request() -> ProcessingRequest: return ProcessingRequest( job_id="arrjob-postgres-001", attempt_no=1, source_file_id="arr_file_ephemeral-001", processor_version="3.0.0", rule_set_sha256="a" * 64, ) class PostgresProcessingStateTests(unittest.TestCase): def test_configuration_and_dispatch_key_are_database_compatible(self) -> None: with patch.dict(os.environ, {}, clear=True): with self.assertRaises(ProcessingError): ProcessingDatabaseConfig.from_environment() with patch.dict( os.environ, {"ARR_DATABASE_URL": "postgresql://synthetic"}, clear=True, ): self.assertEqual( ProcessingDatabaseConfig.from_environment().dsn, "postgresql://synthetic", ) key = processing_dispatch_idempotency_key(request()) self.assertEqual(len(key), 64) self.assertTrue(all(character in "0123456789abcdef" for character in key)) def test_reserve_uses_registered_attempt_as_single_persisted_identity(self) -> None: value = request() key = processing_dispatch_idempotency_key(value) connection = ScriptedConnection( { "FROM ingestion.processing_runs AS run": ( 11, "opera_daily", value.processor_version, value.rule_set_sha256, "queued", key, None, "queued", 12, None, ) } ) state = PostgresProcessingState( ProcessingDatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: connection, ) record = state.reserve(value, value.sha256(), key) self.assertEqual(record.request, value) self.assertEqual(record.idempotency_key, key) self.assertEqual(record.remote_status, "queued") self.assertTrue(connection.closed) def test_start_failure_marks_attempt_and_run_terminal(self) -> None: value = request() key = processing_dispatch_idempotency_key(value) connection = ScriptedConnection( { "FROM ingestion.processing_runs AS run": ( 11, "opera_daily", value.processor_version, value.rule_set_sha256, "queued", key, None, "queued", 12, None, ) } ) state = PostgresProcessingState( ProcessingDatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: connection, ) record = state.fail_start( value.job_id, value.attempt_no, "PROCESSING_REMOTE_REJECTED", ) self.assertEqual(record.remote_status, "failed") self.assertEqual(record.failure_code, "PROCESSING_REMOTE_REJECTED") updates = [ (query, params) for query, params in connection.executed if query.startswith("UPDATE ingestion.processing_") ] self.assertEqual(len(updates), 2) self.assertIn("attempt_status = 'failed'", updates[0][0]) self.assertEqual(updates[0][1], ("PROCESSING_REMOTE_REJECTED", 12)) self.assertIn("run_status = 'failed'", updates[1][0]) self.assertEqual(updates[1][1], ("PROCESSING_REMOTE_REJECTED", 11)) def test_bind_run_casts_nullable_failure_parameters_for_postgres(self) -> None: value = request() key = processing_dispatch_idempotency_key(value) connection = ScriptedConnection( { "FROM ingestion.processing_runs AS run": ( 11, "opera_daily", value.processor_version, value.rule_set_sha256, "queued", key, None, "queued", 12, None, ) } ) state = PostgresProcessingState( ProcessingDatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: connection, ) record = state.bind_run( value.job_id, value.attempt_no, "remote-run-001", "pending", ) self.assertEqual(record.remote_run_id, "remote-run-001") self.assertEqual(record.remote_status, "pending") updates = [ (query, params) for query, params in connection.executed if query.startswith("UPDATE ingestion.processing_") ] self.assertEqual(len(updates), 2) for query, _params in updates: self.assertIn("%s::text IS NULL", query) self.assertIsNone(updates[0][1][2]) self.assertIsNone(updates[0][1][3]) self.assertIsNone(updates[1][1][1]) self.assertIsNone(updates[1][1][2]) def test_source_and_committed_outcome_are_resolved_without_exposing_bytes(self) -> None: source_connection = ScriptedConnection( { "artifact.object_key": ( "arr/jobs/arrjob-postgres-001/attempts/0001/committed/source_xml/source.xml", "source.xml", "b" * 64, 123, "application/xml", ) } ) state = PostgresProcessingState( ProcessingDatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: source_connection, ) source = state.source_for_attempt("arrjob-postgres-001", 1) self.assertEqual(source.role, "source_xml") self.assertEqual(source.sha256, "b" * 64) outcome_connection = ScriptedConnection( { "delivery.delivery_status": ( "committed", "arrjob-postgres-001", date(2026, 7, 27), 19, 3, ) } ) state = PostgresProcessingState( ProcessingDatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: outcome_connection, ) outcome = state.outcome_for_delivery("delivery-arrjob-postgres-001") self.assertIsNotNone(outcome) assert outcome is not None self.assertEqual(outcome.status, "already_committed") self.assertEqual(outcome.daily_version_id, 19) self.assertEqual(outcome.version_no, 3) def test_remote_status_mapping_never_marks_remote_success_as_database_acceptance(self) -> None: self.assertEqual( PostgresProcessingState._database_status("success"), ("delivered", "validating", None, False), ) self.assertEqual( PostgresProcessingState._database_status("failed"), ("failed", "failed", "PROCESSING_REMOTE_FAILED", True), ) with self.assertRaises(ProcessingError): PostgresProcessingState._database_status("unknown") if __name__ == "__main__": unittest.main()