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