116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""Opaque, short-lived, single-read grants for the runtime fetch provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, Optional
|
|
|
|
from arr_ingestion.contracts import OPAQUE_ID_RE, IngestionError
|
|
from arr_storage.store import ManagedObjectStore
|
|
|
|
|
|
MAX_READ_GRANT_TTL_SECONDS = 300
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OpaqueReadGrant:
|
|
handle: str
|
|
job_id: str
|
|
expires_at: datetime
|
|
max_reads: int = 1
|
|
|
|
|
|
@dataclass
|
|
class _GrantRecord:
|
|
public: OpaqueReadGrant
|
|
object_key: str
|
|
consumed: bool = False
|
|
|
|
|
|
class InMemoryReadGrantBroker:
|
|
"""Reference broker for tests; production must use a shared atomic grant ledger."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: ManagedObjectStore,
|
|
*,
|
|
clock: Optional[Callable[[], datetime]] = None,
|
|
) -> None:
|
|
self._store = store
|
|
self._clock = clock or (lambda: datetime.now(timezone.utc))
|
|
self._lock = threading.Lock()
|
|
self._records: Dict[str, _GrantRecord] = {}
|
|
|
|
def issue(
|
|
self,
|
|
*,
|
|
job_id: str,
|
|
object_key: str,
|
|
original_filename: str,
|
|
ttl_seconds: int = MAX_READ_GRANT_TTL_SECONDS,
|
|
) -> OpaqueReadGrant:
|
|
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")
|
|
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()
|
|
public = OpaqueReadGrant(
|
|
handle="arr_file_" + secrets.token_urlsafe(24),
|
|
job_id=job_id,
|
|
expires_at=now + timedelta(seconds=ttl_seconds),
|
|
)
|
|
with self._lock:
|
|
self._records[public.handle] = _GrantRecord(public, object_key)
|
|
return public
|
|
|
|
def materialize(
|
|
self,
|
|
*,
|
|
handle: str,
|
|
job_id: str,
|
|
destination: Path,
|
|
max_bytes: int,
|
|
) -> None:
|
|
now = self._aware_now()
|
|
with self._lock:
|
|
record = self._records.get(handle)
|
|
if (
|
|
record is None
|
|
or record.consumed
|
|
or record.public.job_id != job_id
|
|
or now >= record.public.expires_at
|
|
):
|
|
raise IngestionError("READ_GRANT_INVALID", "read grant is unavailable")
|
|
record.consumed = True
|
|
object_key = record.object_key
|
|
self._store.materialize(object_key, destination, max_bytes)
|
|
|
|
def purge_expired(self) -> int:
|
|
now = self._aware_now()
|
|
with self._lock:
|
|
expired = [
|
|
handle
|
|
for handle, record in self._records.items()
|
|
if record.consumed or now >= record.public.expires_at
|
|
]
|
|
for handle in expired:
|
|
del self._records[handle]
|
|
return len(expired)
|
|
|
|
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)
|