from __future__ import annotations
import tempfile
import unittest
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any
from arr_ingestion.contracts import IngestionError
from arr_storage.filesystem import FilesystemObjectBackend
from arr_storage.postgres_grants import (
PostgresReadGrantBroker,
ReadGrantDatabaseConfig,
)
from arr_storage.store import ManagedObjectStore
from arr_web.run import _parser
NOW = datetime(2026, 7, 28, 12, 0, tzinfo=timezone.utc)
class FakeGrantDatabase:
def __init__(self, *, table_exists: bool = True, database: str = "booking_test") -> None:
self.table_exists = table_exists
self.database = database
self.grants: dict[str, dict[str, Any]] = {}
self.connections: list[FakeGrantConnection] = []
def connect(self, _dsn: str) -> "FakeGrantConnection":
connection = FakeGrantConnection(self)
self.connections.append(connection)
return connection
class FakeGrantConnection:
def __init__(self, database: FakeGrantDatabase) -> None:
self.database = database
self.next_row: tuple[Any, ...] | None = None
self.rowcount = -1
self.closed = False
def transaction(self) -> "FakeGrantConnection":
return self
def cursor(self) -> "FakeGrantConnection":
return self
def __enter__(self) -> "FakeGrantConnection":
return self
def __exit__(self, _kind: object, _value: object, _traceback: object) -> bool:
return False
def execute(self, query: str, params: tuple[Any, ...] | None = None) -> None:
normalized = " ".join(query.split())
self.next_row = None
self.rowcount = -1
if "SELECT current_database()" in normalized:
self.next_row = (self.database.database,)
elif "to_regclass('ingestion.source_read_grants')" in normalized:
self.next_row = (
"ingestion.source_read_grants" if self.database.table_exists else None,
)
elif normalized.startswith("INSERT INTO ingestion.source_read_grants"):
assert params is not None
handle, job_id, attempt_no, object_key, filename, created_at, expires_at = params
self.database.grants[str(handle)] = {
"job_id": str(job_id),
"attempt_no": int(attempt_no),
"object_key": str(object_key),
"filename": str(filename),
"created_at": created_at,
"expires_at": expires_at,
"consumed_at": None,
}
self.rowcount = 1
elif normalized.startswith("UPDATE ingestion.source_read_grants"):
assert params is not None
consumed_at, handle, job_id, current_time = params
grant = self.database.grants.get(str(handle))
if (
grant is not None
and grant["job_id"] == job_id
and grant["consumed_at"] is None
and grant["expires_at"] > current_time
):
grant["consumed_at"] = consumed_at
self.next_row = (grant["object_key"], grant["filename"])
self.rowcount = 1
else:
self.rowcount = 0
elif normalized.startswith("DELETE FROM ingestion.source_read_grants"):
assert params is not None
current_time = params[0]
expired = [
handle
for handle, grant in self.database.grants.items()
if grant["consumed_at"] is not None
or grant["expires_at"] <= current_time
]
for handle in expired:
del self.database.grants[handle]
self.rowcount = len(expired)
def fetchone(self) -> tuple[Any, ...] | None:
return self.next_row
def close(self) -> None:
self.closed = True
class PostgresReadGrantBrokerTests(unittest.TestCase):
def _store(self, root: Path) -> tuple[ManagedObjectStore, str]:
store = ManagedObjectStore(
FilesystemObjectBackend(root / "objects", create=True)
)
source = root / "source.xml"
source.write_bytes(b"|
")
stored = store.upload_committed(
job_id="arrjob-grant-postgres-001",
attempt_no=1,
role="source_xml",
source=source,
original_filename="source.xml",
)
return store, stored.object_key
def test_issue_and_atomic_single_redemption_are_job_bound(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
store, object_key = self._store(root)
database = FakeGrantDatabase()
broker = PostgresReadGrantBroker(
store,
ReadGrantDatabaseConfig("postgresql://synthetic"),
connect=database.connect,
clock=lambda: NOW,
)
broker.assert_ready()
grant = broker.issue(
job_id="arrjob-grant-postgres-001",
object_key=object_key,
original_filename="source.xml",
ttl_seconds=300,
)
self.assertNotIn(object_key, repr(grant))
self.assertEqual(grant.max_reads, 1)
self.assertEqual(
database.grants[grant.handle]["attempt_no"],
1,
)
with self.assertRaises(IngestionError):
broker.materialize(
handle=grant.handle,
job_id="arrjob-other",
destination=root / "wrong.xml",
max_bytes=1024,
)
destination = root / "runtime" / "source.xml"
broker.materialize(
handle=grant.handle,
job_id="arrjob-grant-postgres-001",
destination=destination,
max_bytes=1024,
)
self.assertEqual(destination.read_bytes(), b"|
")
with self.assertRaises(IngestionError):
broker.materialize(
handle=grant.handle,
job_id="arrjob-grant-postgres-001",
destination=root / "second.xml",
max_bytes=1024,
)
self.assertEqual(broker.purge_expired(), 1)
self.assertFalse(database.grants)
self.assertTrue(all(item.closed for item in database.connections))
def test_expiry_ttl_and_missing_schema_fail_closed(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
store, object_key = self._store(root)
current = [NOW]
database = FakeGrantDatabase()
broker = PostgresReadGrantBroker(
store,
ReadGrantDatabaseConfig("postgresql://synthetic"),
connect=database.connect,
clock=lambda: current[0],
)
with self.assertRaises(IngestionError):
broker.issue(
job_id="arrjob-grant-postgres-001",
object_key=object_key,
original_filename="source.xml",
ttl_seconds=301,
)
grant = broker.issue(
job_id="arrjob-grant-postgres-001",
object_key=object_key,
original_filename="source.xml",
ttl_seconds=60,
)
current[0] = NOW + timedelta(seconds=60)
with self.assertRaises(IngestionError):
broker.materialize(
handle=grant.handle,
job_id="arrjob-grant-postgres-001",
destination=root / "expired.xml",
max_bytes=1024,
)
self.assertEqual(broker.purge_expired(), 1)
unavailable = PostgresReadGrantBroker(
store,
ReadGrantDatabaseConfig("postgresql://synthetic"),
connect=FakeGrantDatabase(table_exists=False).connect,
clock=lambda: NOW,
)
with self.assertRaisesRegex(IngestionError, "unavailable"):
unavailable.assert_ready()
def test_target_guard_migration_and_processing_flag(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
store, _object_key = self._store(root)
broker = PostgresReadGrantBroker(
store,
ReadGrantDatabaseConfig("postgresql://synthetic"),
connect=FakeGrantDatabase(database="postgres").connect,
clock=lambda: NOW,
)
with self.assertRaisesRegex(IngestionError, "target"):
broker.assert_ready()
migration = (
Path(__file__).resolve().parents[1]
/ "database"
/ "009_source_read_grants.sql"
).read_text(encoding="utf-8")
self.assertIn("current_database() <> 'booking_test'", migration)
self.assertIn("CREATE TABLE ingestion.source_read_grants", migration)
self.assertIn("interval '5 minutes'", migration)
self.assertIn("consumed_at", migration)
self.assertNotIn("DROP SCHEMA", migration)
self.assertFalse(_parser().parse_args([]).enable_processing)
self.assertTrue(
_parser().parse_args(["--enable-processing"]).enable_processing
)
if __name__ == "__main__":
unittest.main(verbosity=2)