"""Restart-safe PostgreSQL ledger for opaque source XML read grants.""" from __future__ import annotations import os import secrets from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Callable, Optional, Tuple from arr_ingestion.contracts import OPAQUE_ID_RE, IngestionError from arr_storage.grants import MAX_READ_GRANT_TTL_SECONDS, OpaqueReadGrant from arr_storage.store import ManagedObjectStore TARGET_DATABASE = "booking_test" DATABASE_ENV = "ARR_DATABASE_URL" @dataclass(frozen=True) class ReadGrantDatabaseConfig: dsn: str @classmethod def from_environment(cls) -> "ReadGrantDatabaseConfig": dsn = os.environ.get(DATABASE_ENV, "").strip() if not dsn: raise IngestionError( "DATABASE_CONFIG_MISSING", f"{DATABASE_ENV} is required", ) return cls(dsn) def _default_connect(dsn: str) -> Any: try: import psycopg # type: ignore[import-not-found] except ImportError: raise IngestionError( "DATABASE_DRIVER_UNAVAILABLE", "PostgreSQL driver is unavailable", ) from None try: return psycopg.connect(dsn, autocommit=False) except Exception: raise IngestionError( "DATABASE_UNAVAILABLE", "read grant database is unavailable", ) from None class PostgresReadGrantBroker: """Persist and atomically redeem one opaque, job-bound source handle.""" def __init__( self, store: ManagedObjectStore, config: ReadGrantDatabaseConfig, *, connect: Optional[Callable[[str], Any]] = None, clock: Optional[Callable[[], datetime]] = None, ) -> None: self._store = store self._config = config self._connect = connect or _default_connect self._clock = clock or (lambda: datetime.now(timezone.utc)) def assert_ready(self) -> None: def operation(cursor: Any) -> None: cursor.execute("SELECT to_regclass('ingestion.source_read_grants')") row = cursor.fetchone() if row is None or row[0] is None: raise IngestionError( "READ_GRANT_LEDGER_UNAVAILABLE", "source read grant ledger is unavailable", ) self._run(operation, read_only=True) def issue( self, *, job_id: str, object_key: str, original_filename: str, ttl_seconds: int = MAX_READ_GRANT_TTL_SECONDS, ) -> OpaqueReadGrant: self._validate_job_and_ttl(job_id, ttl_seconds) descriptor = self._store.inspect_committed(object_key, original_filename) if descriptor.job_id != job_id or descriptor.role != "source_xml": raise IngestionError("READ_GRANT_INVALID", "read grant target is invalid") now = self._aware_now() grant = OpaqueReadGrant( handle="arr_file_" + secrets.token_urlsafe(24), job_id=job_id, expires_at=now + timedelta(seconds=ttl_seconds), ) def operation(cursor: Any) -> None: cursor.execute( """ INSERT INTO ingestion.source_read_grants ( grant_key, job_key, attempt_no, object_key, original_filename, created_at, expires_at ) VALUES (%s, %s, %s, %s, %s, %s, %s) """, ( grant.handle, job_id, descriptor.attempt_no, descriptor.object_key, descriptor.original_filename, now, grant.expires_at, ), ) self._run(operation) return grant def materialize( self, *, handle: str, job_id: str, destination: Path, max_bytes: int, ) -> None: if ( not isinstance(handle, str) or not OPAQUE_ID_RE.fullmatch(handle) or not isinstance(job_id, str) or not OPAQUE_ID_RE.fullmatch(job_id) or not isinstance(destination, Path) or not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 0 ): raise IngestionError("READ_GRANT_INVALID", "read grant is unavailable") now = self._aware_now() def operation(cursor: Any) -> Tuple[str, str]: cursor.execute( """ UPDATE ingestion.source_read_grants SET consumed_at = %s WHERE grant_key = %s AND job_key = %s AND consumed_at IS NULL AND expires_at > %s RETURNING object_key, original_filename """, (now, handle, job_id, now), ) row = cursor.fetchone() if row is None: raise IngestionError( "READ_GRANT_INVALID", "read grant is unavailable", ) return str(row[0]), str(row[1]) object_key, original_filename = self._run(operation) descriptor = self._store.inspect_committed(object_key, original_filename) if descriptor.job_id != job_id or descriptor.role != "source_xml": raise IngestionError("READ_GRANT_INVALID", "read grant target is invalid") self._store.materialize(object_key, destination, max_bytes) def purge_expired(self) -> int: now = self._aware_now() def operation(cursor: Any) -> int: cursor.execute( """ DELETE FROM ingestion.source_read_grants WHERE consumed_at IS NOT NULL OR expires_at <= %s """, (now,), ) return max(0, int(cursor.rowcount)) return int(self._run(operation)) def _run( self, operation: Callable[[Any], Any], *, read_only: bool = False, ) -> Any: try: connection = self._connect(self._config.dsn) except IngestionError: raise except Exception: raise IngestionError( "DATABASE_UNAVAILABLE", "read grant database is unavailable", ) from None try: with connection.transaction(): with connection.cursor() as cursor: cursor.execute( "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE" + (" READ ONLY" if read_only else "") ) cursor.execute("SET LOCAL statement_timeout = '10s'") cursor.execute("SET LOCAL lock_timeout = '3s'") cursor.execute("SELECT current_database()") row = cursor.fetchone() if row is None or row[0] != TARGET_DATABASE: raise IngestionError( "DATABASE_TARGET_INVALID", "read grant database target is invalid", ) return operation(cursor) except IngestionError: raise except Exception: raise IngestionError( "DATABASE_QUERY_FAILED" if read_only else "DATABASE_WRITE_FAILED", "read grant state could not be read" if read_only else "read grant state could not be persisted", ) from None finally: connection.close() @staticmethod def _validate_job_and_ttl(job_id: str, ttl_seconds: int) -> None: if not isinstance(job_id, str) or not OPAQUE_ID_RE.fullmatch(job_id): raise IngestionError("READ_GRANT_INVALID", "read grant job is invalid") if ( not isinstance(ttl_seconds, int) or isinstance(ttl_seconds, bool) or not 1 <= ttl_seconds <= MAX_READ_GRANT_TTL_SECONDS ): raise IngestionError("READ_GRANT_INVALID", "read grant lifetime is invalid") def _aware_now(self) -> datetime: value = self._clock() if value.tzinfo is None or value.utcoffset() is None: raise ValueError("grant clock must return a timezone-aware datetime") return value.astimezone(timezone.utc)