client-query-cache 0.1.0__py3-none-any.whl

client_query_cache/asynchronous/database.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+
+from pymongo.asynchronous.collection import AsyncCollection
+from pymongo.asynchronous.database import AsyncDatabase
+
+from client_query_cache.asynchronous.collection import CachedCollection
+
+if TYPE_CHECKING:
+    from client_query_cache.asynchronous.manager import CacheManager
+
+
+class CachedDatabase[DocumentType: Mapping[str, Any]]:
+    __slots__ = ("_database", "_manager")
+
+    def __init__(
+        self, manager: CacheManager[DocumentType], database: AsyncDatabase[DocumentType]
+    ) -> None:
+        self._manager = manager
+        self._database = database
+
+    @property
+    def manager(self) -> CacheManager[DocumentType]:
+        return self._manager
+
+    @property
+    def name(self) -> str:
+        return self._database.name
+
+    @property
+    def raw(self) -> AsyncDatabase[DocumentType]:
+        return self._database
+
+    def __getitem__(self, name: str) -> CachedCollection[DocumentType]:
+        return CachedCollection(self, self._database[name])
+
+    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
+        return self._wrap_delegated(getattr(self._database, name))
+
+    def _wrap_delegated(self, value: Any) -> Any:  # noqa: ANN401
+        if isinstance(value, AsyncCollection):
+            return CachedCollection(self, value)
+        if isinstance(value, AsyncDatabase):
+            return CachedDatabase(self._manager, value)
+        if inspect.iscoroutine(value):
+            return self._await_and_wrap(value)
+        if not callable(value):
+            return value
+
+        def _delegate(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            return self._wrap_delegated(value(*args, **kwargs))
+
+        return _delegate
+
+    async def _await_and_wrap(self, coroutine: Any) -> Any:  # noqa: ANN401
+        return self._wrap_delegated(await coroutine)

client_query_cache/asynchronous/manager.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Mapping
+from typing import TYPE_CHECKING, Any, Self
+
+from client_query_cache._core.collection_metadata import (
+    CollectionMetadata,
+    CollectionMetadataCache,
+)
+from client_query_cache._core.manager import CacheCore
+from client_query_cache._core.stream_options import (
+    DEFAULT_MAX_AWAIT_TIME_MS,
+    validate_max_await_time_ms,
+)
+from client_query_cache._core.unique_keys import (
+    UniqueKeyMetadata,
+    UniqueKeyMetadataCache,
+    discover_unique_keys,
+)
+from client_query_cache.asynchronous.database import CachedDatabase
+from client_query_cache.asynchronous.streams import ChangeStreamCoordinator
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from pymongo import AsyncMongoClient
+
+    from client_query_cache._core.collection_metadata import CollectionProbeResult
+    from client_query_cache._core.keys import NamespaceId
+    from client_query_cache._core.manager import CacheCoreConfig
+    from client_query_cache._core.unique_keys import UniqueKeyDefinition
+
+
+class CacheManager[DocumentType: Mapping[str, Any]]:
+    __slots__ = ("_cache", "_client", "_coordinator", "_metadata", "_unique_keys")
+
+    def __init__(
+        self,
+        client: AsyncMongoClient[DocumentType],
+        *,
+        cache_config: CacheCoreConfig | None = None,
+        max_await_time_ms: int = DEFAULT_MAX_AWAIT_TIME_MS,
+    ) -> None:
+        validate_max_await_time_ms(max_await_time_ms)
+        self._client = client
+        self._cache = CacheCore(cache_config)
+        self._coordinator = ChangeStreamCoordinator(
+            client, self._cache, max_await_time_ms=max_await_time_ms
+        )
+        self._metadata = CollectionMetadataCache()
+        self._unique_keys = UniqueKeyMetadataCache()
+
+    @property
+    def client(self) -> AsyncMongoClient[DocumentType]:
+        return self._client
+
+    @property
+    def cache_core(self) -> CacheCore:
+        return self._cache
+
+    async def ensure_cache_eligible(
+        self,
+        namespace: NamespaceId,
+        collection_probe: Callable[[], Awaitable[CollectionProbeResult | None]],
+    ) -> bool:
+        await self._coordinator.activate_database(namespace.database)
+        if not self._cache.is_database_available(namespace.database):
+            return False
+        current_epoch = self._cache.current_epoch(namespace)
+        cached = self._metadata.get(namespace)
+        if cached is None or cached.checked_epoch != current_epoch:
+            probe_result = await collection_probe()
+            if probe_result is None:
+                return False
+            cached = CollectionMetadata(
+                checked_epoch=current_epoch,
+                is_cacheable=probe_result.is_cacheable,
+                default_collation=probe_result.default_collation,
+            )
+            self._metadata.put(namespace, cached)
+        return cached.is_cacheable
+
+    def default_collation_for(self, namespace: NamespaceId) -> Mapping[str, Any] | None:
+        cached = self._metadata.get(namespace)
+        return cached.default_collation if cached is not None else None
+
+    async def unique_keys_for(
+        self,
+        namespace: NamespaceId,
+        list_indexes: Callable[[], Awaitable[Sequence[Mapping[str, Any]] | None]],
+    ) -> tuple[UniqueKeyDefinition, ...]:
+        current_generation = self._cache.current_index_generation(namespace)
+        cached = self._unique_keys.get(namespace)
+        if cached is None or cached.checked_index_generation != current_generation:
+            index_specs = await list_indexes()
+            if index_specs is None:
+                return ()
+            if self._cache.current_index_generation(namespace) != current_generation:
+                return ()
+            cached = UniqueKeyMetadata(
+                checked_index_generation=current_generation,
+                keys=discover_unique_keys(index_specs),
+            )
+            self._unique_keys.put(namespace, cached)
+        return cached.keys
+
+    def __getitem__(self, name: str) -> CachedDatabase[DocumentType]:
+        return CachedDatabase(self, self._client[name])
+
+    async def close(self) -> None:
+        await self._coordinator.close()
+        self._cache.close()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, *_exc_info: object) -> None:
+        await self.close()

client_query_cache/asynchronous/streams.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from typing import TYPE_CHECKING, Any
+
+import bson
+from bson.errors import BSONError
+from pymongo.errors import OperationFailure, PyMongoError
+
+from client_query_cache._core.errors import StreamLifecycleError, StreamStartupError
+from client_query_cache._core.stream_events import (
+    build_change_stream_pipeline,
+    is_unresumable_change_stream_error,
+    route_change_event,
+)
+from client_query_cache._core.stream_health import RetryBackoff, StreamHealth
+from client_query_cache._core.stream_options import DEFAULT_MAX_AWAIT_TIME_MS
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from pymongo import AsyncMongoClient
+    from pymongo.asynchronous.change_stream import AsyncDatabaseChangeStream
+    from pymongo.asynchronous.database import AsyncDatabase
+
+    from client_query_cache._core.manager import CacheCore
+
+MINIMUM_SERVER_VERSION = (8, 0)
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseStreamSupervisor:
+    __slots__ = (
+        "_backoff",
+        "_cache",
+        "_database",
+        "_health",
+        "_max_await_time_ms",
+        "_resume_token",
+        "_stop_event",
+        "_stream",
+        "_task",
+    )
+
+    def __init__(
+        self,
+        database: AsyncDatabase[Any],
+        cache: CacheCore,
+        *,
+        backoff: RetryBackoff | None = None,
+        max_await_time_ms: int = DEFAULT_MAX_AWAIT_TIME_MS,
+    ) -> None:
+        self._database = database
+        self._cache = cache
+        self._backoff = backoff if backoff is not None else RetryBackoff()
+        self._max_await_time_ms = max_await_time_ms
+        self._health = StreamHealth.STARTING
+        self._stream: AsyncDatabaseChangeStream[Any] | None = None
+        self._resume_token: Mapping[str, Any] | None = None
+        self._stop_event = asyncio.Event()
+        self._task: asyncio.Task[None] | None = None
+        self._cache.set_database_available(self._database.name, available=False)
+
+    @property
+    def healthy(self) -> bool:
+        return self._health is StreamHealth.HEALTHY
+
+    async def start(self) -> None:
+        if self._health is not StreamHealth.STARTING:
+            message = "start() may only be called once per supervisor instance"
+            raise StreamLifecycleError(message)
+        self._set_health(StreamHealth.CONNECTING)
+        self._clear_namespaces_for_database()
+        try:
+            await self._ensure_server_supports_expanded_events()
+            await self._open_stream(resume_token=None, use_start_after=False)
+        except StreamStartupError:
+            self._set_health(StreamHealth.CLOSED)
+            raise
+        except PyMongoError as exc:
+            self._set_health(StreamHealth.CLOSED)
+            message = (
+                f"failed to open change stream for database {self._database.name!r}"
+            )
+            raise StreamStartupError(message) from exc
+        except asyncio.CancelledError:
+            self._set_health(StreamHealth.CLOSED)
+            stream = self._stream
+            if stream is not None:
+                with contextlib.suppress(PyMongoError):
+                    await stream.close()
+            raise
+        if self._stop_event.is_set():
+            assert self._stream is not None
+            with contextlib.suppress(PyMongoError):
+                await self._stream.close()
+            self._set_health(StreamHealth.CLOSED)
+            message = "stop() was called while start() was still connecting"
+            raise StreamLifecycleError(message)
+        self._set_health(StreamHealth.HEALTHY)
+        self._task = asyncio.ensure_future(self._run())
+
+    async def stop(self) -> None:
+        self._stop_event.set()
+        self._set_health(StreamHealth.CLOSED)
+        task = self._task
+        if task is not None:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+        stream = self._stream
+        if stream is not None:
+            try:
+                await stream.close()
+            except PyMongoError:
+                logger.warning(
+                    "change stream close failed during shutdown",
+                    extra={"database": self._database.name},
+                    exc_info=True,
+                )
+        self._set_health(StreamHealth.CLOSED)
+
+    def _set_health(self, health: StreamHealth) -> None:
+        if self._stop_event.is_set() and health is StreamHealth.HEALTHY:
+            health = StreamHealth.CLOSED
+        self._cache.set_database_available(
+            self._database.name, available=health is StreamHealth.HEALTHY
+        )
+        self._health = health
+
+    async def _ensure_server_supports_expanded_events(self) -> None:
+        server_info = await self._database.client.server_info()
+        version = tuple(server_info["versionArray"][:2])
+        if version < MINIMUM_SERVER_VERSION:
+            message = (
+                f"MongoDB server version {server_info['version']} is not supported; "
+                "client_query_cache requires MongoDB 8.0 or newer"
+            )
+            raise StreamStartupError(message)
+
+    async def _open_stream(
+        self, *, resume_token: Mapping[str, Any] | None, use_start_after: bool
+    ) -> None:
+        kwargs: dict[str, Any] = {
+            "show_expanded_events": True,
+            "max_await_time_ms": self._max_await_time_ms,
+        }
+        if resume_token is not None:
+            if use_start_after:
+                kwargs["start_after"] = resume_token
+            else:
+                kwargs["resume_after"] = resume_token
+        previous_stream = self._stream
+        self._stream = await self._database.watch(
+            build_change_stream_pipeline(), **kwargs
+        )
+        if previous_stream is not None:
+            with contextlib.suppress(PyMongoError):
+                await previous_stream.close()
+        if self._stop_event.is_set():
+            with contextlib.suppress(PyMongoError):
+                await self._stream.close()
+
+    async def _interruptible_sleep(self, delay: float) -> bool:
+        try:
+            await asyncio.wait_for(self._stop_event.wait(), timeout=delay)
+        except TimeoutError:
+            return False
+        return True
+
+    async def _run(self) -> None:
+        while not self._stop_event.is_set():
+            assert self._stream is not None
+            self._cache.record_stream_poll(self._database.name)
+            try:
+                event = await self._stream.next()
+            except StopAsyncIteration, PyMongoError:
+                await self._handle_stream_failure()
+                continue
+            self._resume_token = self._stream.resume_token
+            must_reopen = route_change_event(self._cache, self._database.name, event)
+            self._record_logical_event_bytes(event)
+            if must_reopen:
+                await self._reopen_after_invalidate()
+
+    def _record_logical_event_bytes(self, event: Mapping[str, Any]) -> None:
+        try:
+            encoded_length = len(
+                bson.encode(dict(event), codec_options=self._database.codec_options)
+            )
+        except BSONError, TypeError, ValueError:
+            return
+        self._cache.record_logical_event_bytes(self._database.name, encoded_length)
+
+    async def _handle_stream_failure(self) -> None:
+        if self._stop_event.is_set():
+            return
+        self._set_health(StreamHealth.RECONNECTING)
+        if self._resume_token is None:
+            self._clear_namespaces_for_database()
+        await self._reopen_with_backoff(use_start_after=False)
+
+    async def _reopen_after_invalidate(self) -> None:
+        self._set_health(StreamHealth.RECONNECTING)
+        await self._reopen_with_backoff(use_start_after=True)
+
+    async def _reopen_with_backoff(self, *, use_start_after: bool) -> None:
+        while not self._stop_event.is_set():
+            try:
+                await self._open_stream(
+                    resume_token=self._resume_token, use_start_after=use_start_after
+                )
+            except PyMongoError as exc:
+                if isinstance(
+                    exc, OperationFailure
+                ) and is_unresumable_change_stream_error(exc):
+                    self._clear_namespaces_for_database()
+                    self._resume_token = None
+                    use_start_after = False
+                    continue
+                delay = self._backoff.next_delay()
+                logger.warning(
+                    "change stream reconnect failed, retrying with backoff",
+                    extra={"database": self._database.name, "delay_seconds": delay},
+                    exc_info=exc,
+                )
+                if await self._interruptible_sleep(delay):
+                    return
+                continue
+            else:
+                self._backoff.reset()
+                if not self._stop_event.is_set():
+                    self._set_health(StreamHealth.HEALTHY)
+                return
+
+    def _clear_namespaces_for_database(self) -> None:
+        for namespace in self._cache.namespaces_for_database(self._database.name):
+            self._cache.clear_namespace(namespace)
+        self._cache.reset_stream_cost_statistics(self._database.name)
+
+
+class ChangeStreamCoordinator:
+    __slots__ = (
+        "_cache",
+        "_client",
+        "_closed",
+        "_lock",
+        "_max_await_time_ms",
+        "_supervisors",
+    )
+
+    def __init__(
+        self,
+        client: AsyncMongoClient[Any],
+        cache: CacheCore,
+        *,
+        max_await_time_ms: int = DEFAULT_MAX_AWAIT_TIME_MS,
+    ) -> None:
+        self._client = client
+        self._max_await_time_ms = max_await_time_ms
+        self._cache = cache
+        self._supervisors: dict[str, DatabaseStreamSupervisor] = {}
+        self._closed = False
+        self._lock = asyncio.Lock()
+
+    async def activate_database(self, name: str) -> DatabaseStreamSupervisor | None:
+        async with self._lock:
+            if self._closed:
+                raise StreamLifecycleError("coordinator is closed")
+            try:
+                supervisor = self._supervisors[name]
+            except KeyError:
+                supervisor = DatabaseStreamSupervisor(
+                    self._client[name],
+                    self._cache,
+                    max_await_time_ms=self._max_await_time_ms,
+                )
+                try:
+                    await supervisor.start()
+                except StreamStartupError:
+                    logger.warning(
+                        "change stream startup failed for database %r; reads for "
+                        "this database will bypass the cache",
+                        name,
+                        exc_info=True,
+                    )
+                    return None
+                self._supervisors[name] = supervisor
+            return supervisor
+
+    async def close(self) -> None:
+        async with self._lock:
+            self._closed = True
+            supervisors = list(self._supervisors.values())
+            self._supervisors.clear()
+        for supervisor in supervisors:
+            await supervisor.stop()

client_query_cache/otel.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+try:
+    from opentelemetry.metrics import Observation
+except ImportError as error:
+    _message = (
+        "OpenTelemetry metrics support requires the 'opentelemetry-api' package; "
+        "install it with the 'client-query-cache[otel]' extra"
+    )
+    raise ImportError(_message) from error
+
+from client_query_cache._core.errors import CacheConfigurationError
+from client_query_cache._core.stream_cost import (
+    INVALIDATION_LAG_CLOCK_SKEW_LIMITATION,
+    RESIDENT_BYTES_SCOPE,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from opentelemetry.metrics import CallbackOptions, Meter
+
+    from client_query_cache._core.manager import CacheCore
+
+__all__ = ["register_cache_metrics"]
+
+_DB_NAMESPACE_ATTRIBUTE = "db.namespace"
+
+_CACHE_COUNTER_FIELDS: tuple[tuple[str, str, str], ...] = (
+    ("client_query_cache.cache.hits", "hits", "Cumulative cache hits."),
+    ("client_query_cache.cache.misses", "misses", "Cumulative cache misses."),
+    (
+        "client_query_cache.cache.evictions",
+        "evictions",
+        "Cumulative least-recently-used cache evictions.",
+    ),
+    (
+        "client_query_cache.cache.bypasses",
+        "bypasses",
+        "Cumulative reads that bypassed the cache.",
+    ),
+    (
+        "client_query_cache.cache.bypasses.oversized",
+        "oversized_bypasses",
+        "Cumulative cache bypasses caused by a result exceeding the max entry size.",
+    ),
+)
+
+_STREAM_COUNTER_FIELDS: tuple[tuple[str, str, str, str], ...] = (
+    (
+        "client_query_cache.stream.polls",
+        "stream_polls",
+        "1",
+        "Cumulative manager calls to change-stream iteration, per database.",
+    ),
+    (
+        "client_query_cache.stream.invalidations",
+        "invalidations",
+        "1",
+        "Cumulative invalidations applied, per database.",
+    ),
+    (
+        "client_query_cache.stream.logical_event_bytes",
+        "logical_event_bytes",
+        "By",
+        "Cumulative logical projected-event bytes observed, per database.",
+    ),
+)
+
+_LAG_GAUGE_DESCRIPTION = (
+    "Simple order-statistic quantile, tagged by the 'percentile' attribute, of "
+    "currently retained invalidation-delivery-lag samples; not a confidence-interval "
+    f"estimate. {INVALIDATION_LAG_CLOCK_SKEW_LIMITATION}"
+)
+
+
+def _percentile(sorted_samples: tuple[float, ...], percentile: float) -> float:
+    index = round(percentile * (len(sorted_samples) - 1))
+    return sorted_samples[index]
+
+
+def _make_cache_counter_callback(
+    cache_core: CacheCore, field: str
+) -> Callable[[CallbackOptions], Iterable[Observation]]:
+    def callback(_options: CallbackOptions) -> Iterable[Observation]:
+        yield Observation(getattr(cache_core.snapshot(), field))
+
+    return callback
+
+
+def _register_cache_counters(meter: Meter, cache_core: CacheCore) -> None:
+    for name, field, description in _CACHE_COUNTER_FIELDS:
+        meter.create_observable_counter(
+            name,
+            callbacks=[_make_cache_counter_callback(cache_core, field)],
+            unit="1",
+            description=description,
+        )
+
+
+def _register_cache_gauges(meter: Meter, cache_core: CacheCore) -> None:
+    def entries_callback(_options: CallbackOptions) -> Iterable[Observation]:
+        yield Observation(cache_core.snapshot().entry_count)
+
+    def resident_bytes_callback(_options: CallbackOptions) -> Iterable[Observation]:
+        yield Observation(cache_core.snapshot().used_bytes)
+
+    meter.create_observable_gauge(
+        "client_query_cache.cache.entries",
+        callbacks=[entries_callback],
+        unit="1",
+        description="Current resident cache entry count.",
+    )
+    meter.create_observable_gauge(
+        "client_query_cache.cache.resident_bytes",
+        callbacks=[resident_bytes_callback],
+        unit="By",
+        description=RESIDENT_BYTES_SCOPE,
+    )
+
+
+def _make_stream_counter_callback(
+    cache_core: CacheCore, field: str
+) -> Callable[[CallbackOptions], Iterable[Observation]]:
+    def callback(_options: CallbackOptions) -> Iterable[Observation]:
+        for database in cache_core.active_stream_cost_databases():
+            snapshot = cache_core.stream_cost_snapshot(database)
+            yield Observation(
+                getattr(snapshot, field), {_DB_NAMESPACE_ATTRIBUTE: database}
+            )
+
+    return callback
+
+
+def _register_stream_counters(meter: Meter, cache_core: CacheCore) -> None:
+    for name, field, unit, description in _STREAM_COUNTER_FIELDS:
+        meter.create_observable_counter(
+            name,
+            callbacks=[_make_stream_counter_callback(cache_core, field)],
+            unit=unit,
+            description=description,
+        )
+
+
+def _make_lag_percentile_callback(
+    cache_core: CacheCore, lag_percentiles: tuple[float, ...]
+) -> Callable[[CallbackOptions], Iterable[Observation]]:
+    def callback(_options: CallbackOptions) -> Iterable[Observation]:
+        for database in cache_core.active_stream_cost_databases():
+            snapshot = cache_core.stream_cost_snapshot(database)
+            samples = tuple(
+                sorted(
+                    value
+                    for window in snapshot.invalidation_lag_windows
+                    for value in window
+                )
+            )
+            if not samples:
+                continue
+            for percentile in lag_percentiles:
+                yield Observation(
+                    _percentile(samples, percentile),
+                    {
+                        _DB_NAMESPACE_ATTRIBUTE: database,
+                        "percentile": percentile,
+                    },
+                )
+
+    return callback
+
+
+def _register_lag_gauges(
+    meter: Meter, cache_core: CacheCore, lag_percentiles: tuple[float, ...]
+) -> None:
+    meter.create_observable_gauge(
+        "client_query_cache.stream.invalidation_lag",
+        callbacks=[_make_lag_percentile_callback(cache_core, lag_percentiles)],
+        unit="s",
+        description=_LAG_GAUGE_DESCRIPTION,
+    )
+
+
+def _validate_lag_percentiles(lag_percentiles: tuple[float, ...]) -> None:
+    for percentile in lag_percentiles:
+        if not math.isfinite(percentile) or not 0.0 <= percentile <= 1.0:
+            message = (
+                f"lag_percentiles values must be finite numbers in [0.0, 1.0]; "
+                f"got {percentile!r}"
+            )
+            raise CacheConfigurationError(message)
+
+
+def register_cache_metrics(
+    meter: Meter,
+    cache_core: CacheCore,
+    *,
+    lag_percentiles: tuple[float, ...] = (0.5, 0.95, 1.0),
+) -> None:
+    _validate_lag_percentiles(lag_percentiles)
+    _register_cache_counters(meter, cache_core)
+    _register_cache_gauges(meter, cache_core)
+    _register_stream_counters(meter, cache_core)
+    _register_lag_gauges(meter, cache_core, lag_percentiles)

client_query_cache/synchronous/__init__.py

@@ -0,0 +1,16 @@
+from client_query_cache._core.manager import CacheCore, CacheCoreConfig
+
+from .collection import CachedCollection
+from .database import CachedDatabase
+from .manager import CacheManager
+from .streams import ChangeStreamCoordinator, DatabaseStreamSupervisor
+
+__all__ = [
+    "CacheCore",
+    "CacheCoreConfig",
+    "CacheManager",
+    "CachedCollection",
+    "CachedDatabase",
+    "ChangeStreamCoordinator",
+    "DatabaseStreamSupervisor",
+]