from __future__ import annotations import os import tempfile import unittest from pathlib import Path from unittest.mock import patch from arr_ingestion.contracts import ArtifactRef, IngestionError from arr_ingestion.postgres import ( DatabaseConfig, PostgresIngestionRepository, ) from arr_ingestion.repository import JobRegistration from tests.test_arr_opera_daily_ingest import run_processor, success_xml class InsertCursor: def __init__(self) -> None: self.params: list[tuple[object, ...]] = [] self._next_id = 100 def execute(self, query: str, params: tuple[object, ...]) -> None: self.assert_placeholder_count(query, params) self.params.append(params) @staticmethod def assert_placeholder_count(query: str, params: tuple[object, ...]) -> None: if query.count("%s") != len(params): raise AssertionError( f"placeholder mismatch: {query.count('%s')} != {len(params)}" ) def fetchone(self) -> tuple[int]: current = self._next_id self._next_id += 1 return (current,) class BookingCursor: def __init__(self) -> None: self.executed: tuple[str, tuple[object, ...]] | None = None def execute(self, query: str, params: tuple[object, ...]) -> None: InsertCursor.assert_placeholder_count(query, params) self.executed = (query, params) def fetchall(self): return [ ("GROUP-MATCH", 1), ("GROUP-REPEATED", 2), ] class ArtifactCursor: def __init__(self, existing=None, inserted_id: int = 101) -> None: self.existing = existing self.inserted_id = inserted_id self.last_query = "" self.executed: list[tuple[str, tuple[object, ...]]] = [] def execute(self, query: str, params: tuple[object, ...]) -> None: InsertCursor.assert_placeholder_count(query, params) self.last_query = " ".join(query.split()) self.executed.append((self.last_query, params)) def fetchone(self): if self.last_query.startswith("SELECT"): return self.existing if self.last_query.startswith("INSERT INTO ingestion.artifacts"): return (self.inserted_id,) return None class SqlStateError(RuntimeError): def __init__(self, sqlstate: str): super().__init__("synthetic database error") self.sqlstate = sqlstate class RetryConnection: def __init__(self) -> 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 close(self) -> None: self.closed = True class LifecycleCursor: def __init__(self, run_status: str = "queued", attempt_status: str = "queued") -> None: self.run_status = run_status self.attempt_status = attempt_status self.last_query = "" self.executed: list[tuple[str, tuple[object, ...]]] = [] def execute(self, query: str, params: tuple[object, ...]) -> None: InsertCursor.assert_placeholder_count(query, params) self.last_query = " ".join(query.split()) self.executed.append((self.last_query, params)) def fetchone(self): if self.last_query.startswith("SELECT run.id, run.run_status"): return (12, self.run_status, 13, self.attempt_status) return None class RegistrationCursor: def __init__(self) -> None: self.last_query = "" self.executed: list[tuple[str, tuple[object, ...]]] = [] def execute(self, query: str, params: tuple[object, ...]) -> None: InsertCursor.assert_placeholder_count(query, params) self.last_query = " ".join(query.split()) self.executed.append((self.last_query, params)) def fetchone(self): if self.last_query.startswith( "SELECT id, artifact_kind, original_filename" ): return None if self.last_query.startswith("INSERT INTO ingestion.artifacts"): return (101,) if self.last_query.startswith( "SELECT id, source_artifact_id, requested_processor_version" ): return None if self.last_query.startswith("INSERT INTO ingestion.processing_runs"): return (202,) return None class PostgresIngestionTests(unittest.TestCase): @staticmethod def artifact(object_key: str, sha256: str = "a" * 64) -> ArtifactRef: return ArtifactRef.from_dict( "source_xml", { "object_key": object_key, "original_filename": "source.xml", "sha256": sha256, "byte_size": 321, "mime_type": "application/xml", }, ) def test_database_config_is_environment_only(self): with patch.dict(os.environ, {}, clear=True): with self.assertRaisesRegex(IngestionError, "ARR_DATABASE_URL"): DatabaseConfig.from_environment() with patch.dict(os.environ, {"ARR_DATABASE_URL": "postgresql://synthetic"}, clear=True): self.assertEqual( DatabaseConfig.from_environment().dsn, "postgresql://synthetic", ) def test_record_insert_maps_duplicate_lineage_and_booking_source_status(self): with tempfile.TemporaryDirectory() as temp_dir: exit_code, _xml, _output, _result, payload = run_processor( success_xml(), Path(temp_dir) ) self.assertEqual(exit_code, 0) cursor = InsertCursor() links = { "AB-123": ("matched", 1), "GO-EASY-GROUP": ("matched", 2), } PostgresIngestionRepository._insert_records( cursor, 9, payload["records"], links, ) self.assertEqual(len(cursor.params), payload["source_rows"]) for params in cursor.params: self.assertEqual(params[0], 9) self.assertEqual(len(params), 33) duplicate_params = cursor.params[2] self.assertEqual(duplicate_params[7], 101) matched_params = cursor.params[1] self.assertEqual(matched_params[-2:], ("matched", 1)) missing_group_params = cursor.params[3] self.assertEqual( missing_group_params[-2:], ("missing_group_code", 0) ) repeated_params = cursor.params[4] self.assertEqual(repeated_params[-2:], ("matched", 2)) def test_booking_lookup_accepts_repeated_group_rows_and_marks_unmatched(self): records = [ {"group_code_key": "GROUP-MATCH"}, {"group_code_key": "GROUP-REPEATED"}, {"group_code_key": "GROUP-MISSING"}, {"group_code_key": None}, ] cursor = BookingCursor() links = PostgresIngestionRepository._load_booking_links(cursor, records) self.assertEqual(links["GROUP-MATCH"], ("matched", 1)) self.assertEqual(links["GROUP-REPEATED"], ("matched", 2)) self.assertEqual(links["GROUP-MISSING"], ("unmatched", 0)) self.assertIsNotNone(cursor.executed) def test_artifact_identity_is_storage_object_not_content_hash(self): key = ( "arr/jobs/arrjob-repeat/attempts/0001/committed/" "source_xml/source.xml" ) cursor = ArtifactCursor() artifact_id = PostgresIngestionRepository._ensure_artifact( cursor, self.artifact(key), ) self.assertEqual(artifact_id, 101) select_query, select_params = cursor.executed[0] self.assertIn("object_key = %s", select_query) self.assertIn("object_version_id IS NULL", select_query) self.assertNotIn("WHERE artifact_kind = %s", select_query) self.assertEqual(select_params, ("oss", "arr-private", key)) self.assertTrue( any(query.startswith("INSERT INTO ingestion.artifacts") for query, _ in cursor.executed) ) def test_existing_storage_identity_is_reused_only_when_all_metadata_matches(self): key = ( "arr/jobs/arrjob-existing/attempts/0001/committed/" "source_xml/source.xml" ) cursor = ArtifactCursor( (77, "opera_xml", "source.xml", "a" * 64, 321, "application/xml") ) artifact_id = PostgresIngestionRepository._ensure_artifact( cursor, self.artifact(key), ) self.assertEqual(artifact_id, 77) self.assertEqual(len(cursor.executed), 1) def test_existing_storage_identity_rejects_conflicting_content(self): key = ( "arr/jobs/arrjob-conflict/attempts/0001/committed/" "source_xml/source.xml" ) cursor = ArtifactCursor( (77, "opera_xml", "source.xml", "b" * 64, 321, "application/xml") ) with self.assertRaises(IngestionError) as raised: PostgresIngestionRepository._ensure_artifact( cursor, self.artifact(key), ) self.assertEqual(raised.exception.code, "ARTIFACT_CONFLICT") def test_job_registration_persists_uploaded_filename_separately(self): cursor = RegistrationCursor() registration = JobRegistration( job_id="arrjob-upload-name-001", source=self.artifact( "arr/jobs/arrjob-upload-name-001/attempts/0001/" "committed/source_xml/source.xml" ), processor_version="3.0.0", rule_set_sha256="b" * 64, attempt_no=1, idempotency_key="c" * 64, uploaded_filename="0723-Opera.XML", ) repository = PostgresIngestionRepository( DatabaseConfig("postgresql://synthetic") ) repository._register_job(cursor, registration) run_insert = next( (query, params) for query, params in cursor.executed if query.startswith("INSERT INTO ingestion.processing_runs") ) self.assertIn("uploaded_filename", run_insert[0]) self.assertEqual(run_insert[1][-1], "0723-Opera.XML") artifact_insert = next( (query, params) for query, params in cursor.executed if query.startswith("INSERT INTO ingestion.artifacts") ) self.assertEqual(artifact_insert[1][4], "source.xml") def test_transaction_reopens_after_serialization_failure(self): connections = [RetryConnection(), RetryConnection()] connect_count = 0 operation_count = 0 def connect(_dsn: str): nonlocal connect_count connection = connections[connect_count] connect_count += 1 return connection def operation(_cursor): nonlocal operation_count operation_count += 1 if operation_count == 1: raise SqlStateError("40001") return "committed" repository = PostgresIngestionRepository( DatabaseConfig("postgresql://synthetic"), connect=connect ) with patch.object(repository, "_begin"), patch("arr_ingestion.postgres.time.sleep"): result = repository._run_transaction(operation, "safe failure") self.assertEqual(result, "committed") self.assertEqual(connect_count, 2) self.assertEqual(operation_count, 2) self.assertTrue(all(connection.closed for connection in connections)) def test_nontransient_database_failure_is_not_retried(self): connection = RetryConnection() repository = PostgresIngestionRepository( DatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: connection ) with patch.object(repository, "_begin"): with self.assertRaisesRegex(IngestionError, "safe failure") as raised: repository._run_transaction( lambda _cursor: (_ for _ in ()).throw(SqlStateError("22000")), "safe failure", ) self.assertEqual(raised.exception.code, "DATABASE_WRITE_FAILED") self.assertTrue(connection.closed) def test_programmatic_start_marks_attempt_and_run_running(self): cursor = LifecycleCursor() repository = PostgresIngestionRepository(DatabaseConfig("postgresql://synthetic")) repository._mark_running( cursor, "arrjob-programmatic-001", 1, ) statements = "\n".join(query for query, _params in cursor.executed) self.assertIn("attempt_status = 'running'", statements) self.assertIn("run_status = 'running'", statements) self.assertIn("started_at = COALESCE(started_at, now())", statements) def test_programmatic_infrastructure_failure_is_terminal_and_outboxed(self): cursor = LifecycleCursor(run_status="running", attempt_status="running") repository = PostgresIngestionRepository(DatabaseConfig("postgresql://synthetic")) repository._record_runtime_failure( cursor, "arrjob-programmatic-002", 1, "PROCESSOR_TIMEOUT", ) statements = "\n".join(query for query, _params in cursor.executed) self.assertIn("attempt_status = 'failed'", statements) self.assertIn("run_status = 'failed'", statements) self.assertIn("INSERT INTO ingestion.outbox_events", statements) rendered_params = repr([params for _query, params in cursor.executed]) self.assertIn("PROCESSOR_TIMEOUT", rendered_params) self.assertNotIn("password", rendered_params.lower()) class Migration008ContractTests(unittest.TestCase): def test_rebuild_targets_v3_and_all_source_record_storage(self): project_root = Path(__file__).resolve().parents[1] sql = (project_root / "database" / "008_arr_mvp_v1_rebuild.sql").read_text( encoding="utf-8" ) for required in ( "CREATE TABLE ingestion.processing_runs", "CREATE TABLE ingestion.processing_attempts", "CREATE TABLE ingestion.processing_deliveries", "CREATE TABLE ingestion.outbox_events", "structured_result_artifact_id", "CREATE TABLE booking.source_rows", "CREATE TABLE booking.parse_versions", "CREATE TABLE finance.daily_records", "CREATE VIEW finance.v_monthly_report_rows", "CREATE VIEW finance.v_channel_details", "outcome <> 'retained'", "total_price = real_price * no_of_rooms * nights", ): with self.subTest(required=required): self.assertIn(required, sql) self.assertNotIn("CREATE TABLE finance.report_versions", sql) self.assertNotIn("CREATE TABLE booking.file_objects", sql) if __name__ == "__main__": unittest.main(verbosity=2)