255 lines
8.6 KiB
Python
255 lines
8.6 KiB
Python
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()
|