lorebase-keeper 0.2.3__py3-none-any.whl

lorebase_keeper/readiness.py

@@ -0,0 +1,257 @@
+"""Readiness checks behind `/readyz`.
+
+A check is an async callable with a `name`. The database check connects, runs `SELECT 1` and
+compares the Alembic heads the code expects (`packaged_heads`, from the store) with the revision
+stored in the `alembic_version` table, so a service whose migrations are missing or behind is not
+ready. Check failures are logged in full and reported to callers as a short
+reason that never contains connection strings or credentials.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import importlib
+import logging
+import re
+from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Any, Protocol
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from lorebase_platform.secrets import SecretResolutionError, resolve_secret
+from lorebase_platform.settings import ConfigError, Settings
+from lorebase_store import bm25
+from lorebase_store.migrations import head_revisions
+from lorebase_store.replicas import ReplicaState
+
+logger = logging.getLogger("lorebase.keeper.readiness")
+
+MigrationHeads = Callable[[], Collection[str]]
+"""Returns the Alembic head revisions the running code expects; empty skips the comparison."""
+
+_DRIVER_SUFFIX = re.compile(r"^(postgres(?:ql)?)\+\w+://")
+
+
+@functools.cache
+def packaged_heads() -> tuple[str, ...]:
+    """Alembic head revisions of the installed store code (read once: they cannot change)."""
+    return head_revisions()
+
+
+class CheckStatus(StrEnum):
+    """Outcome of one readiness check."""
+
+    OK = "ok"
+    DEGRADED = "degraded"
+    NOT_CONFIGURED = "not_configured"
+    FAILED = "failed"
+
+
+@dataclass(frozen=True)
+class CheckResult:
+    """Status of a check and an optional short reason."""
+
+    status: CheckStatus
+    detail: str | None = None
+
+    @property
+    def passed(self) -> bool:
+        """Failed checks make the service not ready; degraded and unconfigured ones do not."""
+        return self.status is not CheckStatus.FAILED
+
+
+class ReadinessCheck(Protocol):
+    """One probe of a dependency."""
+
+    name: str
+
+    async def __call__(self) -> CheckResult:
+        """Run the probe."""
+        ...
+
+
+class NotConfiguredCheck:
+    """Placeholder reporting that a dependency is not part of this deployment."""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    async def __call__(self) -> CheckResult:
+        return CheckResult(CheckStatus.NOT_CONFIGURED)
+
+
+def asyncpg_dsn(dsn: str) -> str:
+    """Connection string without a SQLAlchemy driver suffix (`postgresql+asyncpg://` etc.)."""
+    return _DRIVER_SUFFIX.sub(r"\1://", dsn)
+
+
+def _asyncpg_connect(dsn: str, timeout: float) -> Awaitable[Any]:
+    asyncpg = importlib.import_module("asyncpg")
+    return asyncpg.connect(dsn, timeout=timeout)  # type: ignore[no-any-return]
+
+
+class DatabaseCheck:
+    """Postgres is reachable and its migrations match the code."""
+
+    name = "database"
+
+    def __init__(
+        self,
+        dsn: str,
+        *,
+        timeout_s: float,
+        expected_heads: MigrationHeads | None = None,
+        connect: Callable[[str, float], Awaitable[Any]] = _asyncpg_connect,
+    ) -> None:
+        self._dsn = asyncpg_dsn(dsn)
+        self._timeout_s = timeout_s
+        self._expected_heads = expected_heads
+        self._connect = connect
+
+    @property
+    def expected_heads(self) -> frozenset[str]:
+        """Revisions the database must be at; empty when migrations are not verified."""
+        return frozenset(self._expected_heads()) if self._expected_heads else frozenset()
+
+    async def __call__(self) -> CheckResult:
+        connection = await self._connect(self._dsn, self._timeout_s)
+        try:
+            await connection.fetchval("SELECT 1")
+            expected = self.expected_heads
+            if not expected:
+                return CheckResult(CheckStatus.OK, "migration revision not verified")
+            if await connection.fetchval("SELECT to_regclass('alembic_version')") is None:
+                return CheckResult(CheckStatus.FAILED, "migrations have not been applied")
+            rows = await connection.fetch("SELECT version_num FROM alembic_version")
+        finally:
+            await connection.close()
+        current = frozenset(row["version_num"] for row in rows)
+        if current != expected:
+            return CheckResult(
+                CheckStatus.FAILED,
+                f"migrations at {_join(current)}, expected {_join(expected)}",
+            )
+        return CheckResult(CheckStatus.OK)
+
+
+class ReplicaProbe(Protocol):
+    """What the replica check needs from the store's replica set."""
+
+    async def probe(self, timeout_s: float) -> list[ReplicaState]:
+        """Probe every replica."""
+        ...
+
+
+class ReplicaCheck:
+    """State of the read replicas.
+
+    Replicas that are unreachable or lag more than `max_lag_ms` degrade the service without
+    taking it out of rotation, since reads fall back to the primary, unless `required`: reads
+    are served by replicas only (`prefer_replica` with no fallback) and none is usable. The
+    detail names each replica by its configuration name and `host:port`, never a credential.
+    """
+
+    name = "database_replicas"
+
+    def __init__(
+        self,
+        replicas: ReplicaProbe,
+        *,
+        timeout_s: float,
+        max_lag_ms: int,
+        required: bool,
+    ) -> None:
+        self._replicas = replicas
+        self._timeout_s = timeout_s
+        self._max_lag_ms = max_lag_ms
+        self._required = required
+
+    async def __call__(self) -> CheckResult:
+        states = await self._replicas.probe(self._timeout_s)
+        usable = [s for s in states if s.usable(self._max_lag_ms)]
+        detail = "; ".join(self._describe(s) for s in states)
+        if len(usable) == len(states):
+            return CheckResult(CheckStatus.OK, detail)
+        if not usable and self._required:
+            return CheckResult(CheckStatus.FAILED, detail)
+        return CheckResult(CheckStatus.DEGRADED, detail)
+
+    def _describe(self, state: ReplicaState) -> str:
+        where = f"{state.name} ({state.endpoint})"
+        if not state.reachable:
+            return f"{where}: unreachable ({state.error})"
+        if state.lagging(self._max_lag_ms):
+            lag = "unknown" if state.lag_ms is None else f"{state.lag_ms:.0f} ms"
+            return f"{where}: lagging ({lag})"
+        if state.in_recovery is False:
+            return f"{where}: ok, not a standby"
+        return f"{where}: ok, lag {state.lag_ms or 0:.0f} ms"
+
+
+class SearchBackendCheck:
+    """Whether `search.backend: bm25` is served: degraded, without taking the service out of
+    rotation, while the database has no `pg_search` or no valid BM25 index, when lexical
+    search ranks the default way (`lorebase_store.bm25`)."""
+
+    name = "search_backend"
+
+    def __init__(self, sessions: async_sessionmaker[AsyncSession], *, definer: bool) -> None:
+        self._sessions = sessions
+        self._definer = definer
+
+    async def __call__(self) -> CheckResult:
+        if not self._definer:
+            reason = "store.reads.use_definer_functions is false"
+        else:
+            async with self._sessions() as session:
+                found = await bm25.status(session)
+            if found.ready:
+                return CheckResult(CheckStatus.OK, "bm25")
+            reason = str(found.reason)
+        return CheckResult(
+            CheckStatus.DEGRADED,
+            f"search.backend is bm25 but {reason}; lexical search ranks with postgres",
+        )
+
+
+def _join(revisions: Collection[str]) -> str:
+    return ", ".join(sorted(revisions)) or "none"
+
+
+async def run_checks(checks: Sequence[ReadinessCheck], timeout_s: float) -> dict[str, CheckResult]:
+    """Run every check concurrently, each within `timeout_s`."""
+
+    async def run(check: ReadinessCheck) -> CheckResult:
+        try:
+            return await asyncio.wait_for(check(), timeout_s)
+        except TimeoutError:
+            return CheckResult(CheckStatus.FAILED, f"no answer within {timeout_s:g}s")
+        except Exception as exc:
+            logger.warning("readiness check %s failed", check.name, exc_info=True)
+            return CheckResult(CheckStatus.FAILED, type(exc).__name__)
+
+    results = await asyncio.gather(*(run(check) for check in checks))
+    return {check.name: result for check, result in zip(checks, results, strict=True)}
+
+
+def default_checks(
+    settings: Settings,
+    *,
+    section: str,
+    dsn_key: str,
+    timeout_s: float,
+    expected_heads: MigrationHeads | None,
+) -> list[ReadinessCheck]:
+    """The database check when the settings hold a connection string, else a placeholder."""
+    raw: Mapping[str, Any] = (settings.model_extra or {}).get(section) or {}
+    dsn = raw.get(dsn_key)
+    if not dsn:
+        return [NotConfiguredCheck("database")]
+    try:
+        resolved = resolve_secret(str(dsn))
+    except SecretResolutionError as exc:
+        raise ConfigError([f"{section}.{dsn_key}: {exc}"]) from None
+    return [DatabaseCheck(resolved, timeout_s=timeout_s, expected_heads=expected_heads)]

lorebase_keeper/reading/__init__.py

@@ -0,0 +1 @@
+"""The read path of Keeper (design 8): Records, Links, queries, search, history and diff."""

lorebase_keeper/reading/caching.py

@@ -0,0 +1,251 @@
+"""Conditional GET and response caching of reads.
+
+Every GET read answers with `ETag` = the Revision it was served from plus a digest of the path
+and the canonical query parameters, `Cache-Control: private, must-revalidate` and the
+`Lorebase-Revision` header; `If-None-Match` with that tag answers 304. Responses are cached
+behind `ResponseCache`, keyed by `(Domain, Revision, path, canonical query)`. Revisions are
+immutable, so a commit never invalidates an entry: it moves the branch to a new Revision and
+therefore to new keys. The cache is consulted after the read is resolved and authorized, so it
+never serves a caller something it could not read. Bodies are masked for the caller's
+clearance, which is part of the key (`variant`), and hold vault references, not vault values:
+`disclose` replaces them on every answer, cached or not (`reading.sensitivity`).
+
+Only the in-process `InMemoryResponseCache` exists; a shared adapter (for example over Valkey)
+implements the same protocol.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from collections import OrderedDict
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import Protocol
+
+from fastapi import Request, Response
+from opentelemetry import metrics
+from pydantic import BaseModel
+
+from lorebase_keeper.auth.read_audit import note_all, noted_sensitive
+from lorebase_keeper.settings import ReadCacheSettings
+from lorebase_keeper.tracing.notes import note_cache
+
+JSON = "application/json"
+
+Disclose = Callable[[bytes, str], Awaitable[bytes]]
+"""Rewrites a body of a media type as it is sent (vault references, `reading.sensitivity`)."""
+CACHE_CONTROL = "private, must-revalidate"
+
+_LOOKUPS = metrics.get_meter("lorebase.keeper").create_counter(
+    "lorebase.keeper.cache.requests",
+    unit="{request}",
+    description="Read response cache lookups, by result (hit, miss, bypass when disabled).",
+)
+
+
+@dataclass(frozen=True, slots=True)
+class CacheKey:
+    domain_id: str
+    revision: str
+    path: str
+    query: tuple[tuple[str, str], ...]
+    variant: str = ""
+    """What else the body depends on, such as the Kinds the caller may read."""
+
+
+@dataclass(frozen=True, slots=True)
+class CachedResponse:
+    body: bytes
+    headers: Mapping[str, str] = field(default_factory=dict)
+    time_travel: bool = False
+    media_type: str = JSON
+    sensitive: frozenset[str] = frozenset()
+    """Sensitive values the first rendering showed or masked (`read_audit.note_sensitive`),
+    reported again by every answer from the cache."""
+
+
+class ResponseCache(Protocol):
+    """Storage of read responses (in process now; a shared store later)."""
+
+    async def get(self, key: CacheKey) -> CachedResponse | None: ...
+
+    async def put(self, key: CacheKey, value: CachedResponse) -> None: ...
+
+
+class InMemoryResponseCache:
+    """LRU bounded by entries and bytes, with a time to live and a cap on time travel entries."""
+
+    def __init__(
+        self,
+        settings: ReadCacheSettings,
+        *,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._settings = settings
+        self._clock = clock
+        self._entries: OrderedDict[CacheKey, tuple[float, CachedResponse]] = OrderedDict()
+        self._bytes = 0
+        self._time_travel = 0
+        self._lock = Lock()
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    @property
+    def size_bytes(self) -> int:
+        return self._bytes
+
+    async def get(self, key: CacheKey) -> CachedResponse | None:
+        with self._lock:
+            found = self._entries.get(key)
+            if found is None:
+                return None
+            expires, value = found
+            if expires <= self._clock():
+                self._remove(key)
+                return None
+            self._entries.move_to_end(key)
+            return value
+
+    async def put(self, key: CacheKey, value: CachedResponse) -> None:
+        settings = self._settings
+        size = len(value.body)
+        if size > settings.max_body_bytes or size > settings.max_bytes:
+            return
+        with self._lock:
+            if key in self._entries:
+                self._remove(key)
+            if value.time_travel and self._time_travel >= settings.time_travel_max_entries:
+                return
+            self._entries[key] = (self._clock() + settings.ttl_s, value)
+            self._bytes += size
+            self._time_travel += value.time_travel
+            while len(self._entries) > settings.max_entries or self._bytes > settings.max_bytes:
+                self._remove(next(iter(self._entries)))
+
+    def _remove(self, key: CacheKey) -> None:
+        _, value = self._entries.pop(key)
+        self._bytes -= len(value.body)
+        self._time_travel -= value.time_travel
+
+
+def build_cache(settings: ReadCacheSettings) -> ResponseCache | None:
+    return InMemoryResponseCache(settings) if settings.enabled else None
+
+
+def cache_key(request: Request, domain_id: str, revision: str, variant: str = "") -> CacheKey:
+    query = tuple(sorted(request.query_params.multi_items()))
+    return CacheKey(domain_id, revision, request.url.path, query, variant)
+
+
+def entity_tag(key: CacheKey) -> str:
+    """Strong ETag: the Revision plus a digest of the path and canonical query parameters."""
+    digest = hashlib.sha256(key.path.encode())
+    for name, value in key.query:
+        digest.update(b"\x00" + name.encode() + b"\x1f" + value.encode())
+    if key.variant:
+        digest.update(b"\x1e" + key.variant.encode())
+    return f'"{key.revision}.{digest.hexdigest()[:16]}"'
+
+
+def _matches(request: Request, tag: str) -> bool:
+    header = request.headers.get("if-none-match")
+    if header is None:
+        return False
+    candidates = {t.strip() for t in header.split(",")}
+    return "*" in candidates or tag in candidates or f"W/{tag}" in candidates
+
+
+def not_modified(
+    request: Request,
+    *,
+    domain_id: str,
+    revision: str,
+    headers: Mapping[str, str] | None = None,
+    variant: str = "",
+) -> Response | None:
+    """The 304 answer of a GET read whose `If-None-Match` names its current tag, else `None`."""
+    tag = entity_tag(cache_key(request, domain_id, revision, variant))
+    if not _matches(request, tag):
+        return None
+    common = {"ETag": tag, "Cache-Control": CACHE_CONTROL, **(headers or {})}
+    note_cache("not_modified")
+    return Response(status_code=304, headers=common)
+
+
+async def respond(
+    request: Request,
+    cache: ResponseCache | None,
+    *,
+    domain_id: str,
+    revision: str,
+    produce: Callable[[], Awaitable[BaseModel]],
+    time_travel: bool = False,
+    headers: Mapping[str, str] | None = None,
+    variant: str = "",
+    disclose: Disclose | None = None,
+) -> Response:
+    """Answer a GET read: 304, a cached body, or `produce()` serialized and cached.
+
+    `revision` names everything the response depends on besides the request (one Revision id,
+    or several joined for a diff); `variant` anything else it depends on, such as the Kinds
+    the caller may read and its clearance; `headers` are added to every answer; `disclose`
+    post-processes every body sent (not the cached one).
+    """
+
+    async def body() -> tuple[bytes, str]:
+        return (await produce()).model_dump_json().encode(), JSON
+
+    return await respond_with(
+        request,
+        cache,
+        cache_key(request, domain_id, revision, variant),
+        body,
+        time_travel=time_travel,
+        headers=headers,
+        disclose=disclose,
+    )
+
+
+async def respond_with(
+    request: Request,
+    cache: ResponseCache | None,
+    key: CacheKey,
+    produce: Callable[[], Awaitable[tuple[bytes, str]]],
+    *,
+    time_travel: bool = False,
+    headers: Mapping[str, str] | None = None,
+    disclose: Disclose | None = None,
+) -> Response:
+    """Answer a read under `key`: 304, a cached body, or `produce()` (body and media type),
+    cached. The ETag is `entity_tag(key)`. `disclose`, when given, rewrites bodies as they are
+    sent."""
+    tag = entity_tag(key)
+    common = {"ETag": tag, "Cache-Control": CACHE_CONTROL, **(headers or {})}
+    if _matches(request, tag):
+        note_cache("not_modified")
+        return Response(status_code=304, headers=common)
+    cached = None if cache is None else await cache.get(key)
+    result = "bypass" if cache is None else "miss" if cached is None else "hit"
+    _LOOKUPS.add(1, {"result": result})
+    note_cache(result)
+    if cached is not None:
+        note_all(cached.sensitive)
+        return Response(
+            await _sent(cached.body, cached.media_type, disclose),
+            media_type=cached.media_type,
+            headers={**cached.headers, **common},
+        )
+    body, media_type = await produce()
+    if cache is not None:
+        await cache.put(
+            key,
+            CachedResponse(body, dict(headers or {}), time_travel, media_type, noted_sensitive()),
+        )
+    return Response(await _sent(body, media_type, disclose), media_type=media_type, headers=common)
+
+
+async def _sent(body: bytes, media_type: str, disclose: Disclose | None) -> bytes:
+    return body if disclose is None else await disclose(body, media_type)