252 lines
8.5 KiB
Python
252 lines
8.5 KiB
Python
"""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)
|