feat: prepare ARR for controlled public deployment
This commit is contained in:
Wyndham ARR
253
tests/test_arr_postgres_grants.py
Normal file
253
tests/test_arr_postgres_grants.py
Normal file
@@ -0,0 +1,253 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user