254 lines
9.5 KiB
Python
254 lines
9.5 KiB
Python
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"<root><row /></root>")
|
|
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"<root><row /></root>")
|
|
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)
|