sqlacache 0.1.0__py3-none-any.whl

sqlacache/manager.py

@@ -0,0 +1,348 @@
+"""Cache manager primitives."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from sqlalchemy import event, select
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
+from sqlalchemy.orm import Session
+
+from sqlacache.interceptor import (
+    build_bulk_delete_handler,
+    build_bulk_update_handler,
+    build_do_orm_execute_handler,
+    build_invalidation_handler,
+    handle_bulk_mutation,
+    merge_cached_result,
+    resolve_cached_result,
+)
+from sqlacache.invalidation import _bump_table_version, _get_table_version, generate_tags, invalidate_tags
+from sqlacache.transport.cashews import CashewsTransport
+from sqlacache.utils.key_generation import generate_cache_key
+from sqlacache.utils.query_analysis import extract_model_from_statement, extract_pk_from_instance
+
+if TYPE_CHECKING:
+    from sqlacache.pubsub.redis import RedisPubSub
+    from sqlacache.transport import CacheTransport
+
+
+class CacheManager:
+    """Manage normalized cache configuration and runtime bindings."""
+
+    _async_get_patched = False
+    _original_async_get: Any = None
+    _engine_registry: ClassVar[dict[int, CacheManager]] = {}
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        self._config = config
+        self._transport: CacheTransport | None = None
+        self._model_config: dict[str, dict[str, Any] | None] = dict(config.get("models", {}))
+        self._wildcard_config: dict[str, Any] | None = config.get("wildcard")
+        self._bound_engine: Any | None = None
+        self._bound_sync_engine: Any | None = None
+        self._pubsub_task: asyncio.Task[Any] | None = None
+        self._pubsub: RedisPubSub | None = None
+        self._listeners: list[tuple[Any, str, Any]] = []
+        self._pending_tasks: set[asyncio.Task[Any]] = set()
+
+    @property
+    def config(self) -> dict[str, Any]:
+        return self._config
+
+    async def bind(self, engine: Any) -> None:
+        if self._transport is not None:
+            await self.disconnect()
+        self._bound_engine = engine
+        self._bound_sync_engine = engine.sync_engine if isinstance(engine, AsyncEngine) else engine
+        if self._bound_sync_engine is not None:
+            self._engine_registry[id(self._bound_sync_engine)] = self
+        await self._ensure_transport()
+        self._patch_async_get()
+        self._register_listeners()
+        await self._maybe_setup_pubsub()
+
+    def get_model_config(self, model: type[Any]) -> dict[str, Any] | None:
+        model_path = f"{model.__module__}.{model.__name__}"
+        if model_path in self._model_config:
+            return self._model_config[model_path]
+        return self._wildcard_config
+
+    def is_enabled(self, model: type[Any], op: str) -> bool:
+        model_config = self.get_model_config(model)
+        if not model_config:
+            return False
+        return op in model_config["ops"]
+
+    async def execute(self, session: Any, statement: Any, timeout: int | None = None) -> Any:
+        await self._ensure_transport()
+        models = self._extract_models(statement)
+        execution_statement = statement.execution_options(sqlacache_skip_interceptor=True)
+        if not models:
+            return await session.execute(execution_statement)
+
+        key = await self._build_cache_key(statement, models)
+        transport = self._transport
+        assert transport is not None
+        cached = await transport.get(key)
+        if cached is not None:
+            return merge_cached_result(session, statement, cached)
+
+        result = await session.execute(execution_statement)
+        frozen = result.freeze()
+        primary_model = models[0]
+        model_config = self.get_model_config(primary_model)
+        expire = timeout or (model_config["timeout"] if model_config else self._config["default_timeout"])
+        from sqlacache.utils.query_analysis import extract_pks_from_fetch_result
+
+        pks_by_model = extract_pks_from_fetch_result(list(frozen.data), models)
+        tags: list[str] = []
+        for model, pks in pks_by_model.items():
+            tags.extend(generate_tags(model, pks))
+        await transport.set(key, frozen, expire=expire, tags=tags)
+        return merge_cached_result(session, statement, frozen)
+
+    async def invalidate(self, model: type[Any] | None = None, pks: list[Any] | None = None) -> None:
+        await self._ensure_transport()
+        if model is None:
+            await self.invalidate_all()
+            return
+        transport = self._transport
+        assert transport is not None
+        if pks:
+            await invalidate_tags(transport, *generate_tags(model, pks))
+            await self._publish_invalidation(model, pks, action="manual")
+            return
+        await _bump_table_version(transport, model)
+        await self._publish_invalidation(model, [], action="table")
+
+    async def invalidate_all(self) -> None:
+        await self._ensure_transport()
+        assert self._transport is not None
+        await self._transport.clear()
+
+    async def disconnect(self) -> None:
+        for target, identifier, listener in self._listeners:
+            with suppress(Exception):
+                event.remove(target, identifier, listener)
+        self._listeners.clear()
+        if self._pubsub_task is not None:
+            self._pubsub_task.cancel()
+            self._pubsub_task = None
+        if self._pubsub is not None:
+            await self._pubsub.disconnect()
+            self._pubsub = None
+        for task in list(self._pending_tasks):
+            task.cancel()
+        self._pending_tasks.clear()
+        if self._transport is not None:
+            await self._transport.disconnect()
+            self._transport = None
+        if self._bound_sync_engine is not None:
+            self._engine_registry.pop(id(self._bound_sync_engine), None)
+        if not self._engine_registry:
+            self._restore_async_get()
+        self._bound_engine = None
+        self._bound_sync_engine = None
+
+    def bind_sync(self, engine: Any) -> None:
+        raise NotImplementedError("Sync session support is deferred to v0.2.0")
+
+    def execute_sync(self, session: Any, statement: Any, timeout: int | None = None) -> Any:
+        return self._run_sync(self.execute(session, statement, timeout=timeout))
+
+    def invalidate_sync(self, model: type[Any] | None = None, pks: list[Any] | None = None) -> None:
+        self._run_sync(self.invalidate(model=model, pks=pks))
+
+    def _matches_session(self, session: Session) -> bool:
+        if self._bound_sync_engine is None:
+            return False
+        try:
+            bind = session.get_bind()
+        except Exception:
+            return False
+        return bind is self._bound_sync_engine
+
+    def _extract_models(self, statement: Any) -> list[type[Any]]:
+        extracted = extract_model_from_statement(statement)
+        if extracted is None:
+            return []
+        if isinstance(extracted, list):
+            return extracted
+        return [extracted]
+
+    async def _build_cache_key(self, statement: Any, models: list[type[Any]]) -> str:
+        prefix = self._config["prefix"]
+        base_key = generate_cache_key(statement, prefix=prefix)
+        if not models:
+            return base_key
+        transport = self._transport
+        assert transport is not None
+        versions = [str(await _get_table_version(transport, model)) for model in models]
+        return f"{base_key}:v{'.'.join(versions)}"
+
+    def _register_listeners(self) -> None:
+        from sqlacache.config import _resolve_model
+
+        select_listener = build_do_orm_execute_handler(self)
+        self._listeners.append((Session, "do_orm_execute", select_listener))
+        event.listen(Session, "do_orm_execute", select_listener, retval=True)
+
+        for model_path in self._model_config:
+            if model_path == "*":
+                continue
+            try:
+                model = _resolve_model(model_path)
+            except Exception:
+                continue
+            for identifier, action in (
+                ("after_insert", "insert"),
+                ("after_update", "update"),
+                ("after_delete", "delete"),
+            ):
+                listener = build_invalidation_handler(self, action)
+                self._listeners.append((model, identifier, listener))
+                event.listen(model, identifier, listener, propagate=True)
+
+        bulk_update_listener = build_bulk_update_handler(self)
+        self._listeners.append((Session, "after_bulk_update", bulk_update_listener))
+        event.listen(Session, "after_bulk_update", bulk_update_listener)
+
+        bulk_delete_listener = build_bulk_delete_handler(self)
+        self._listeners.append((Session, "after_bulk_delete", bulk_delete_listener))
+        event.listen(Session, "after_bulk_delete", bulk_delete_listener)
+
+    async def _maybe_setup_pubsub(self) -> None:
+        from sqlacache.pubsub.redis import RedisPubSub
+
+        backend = self._config["backend"]
+        redis_url = backend["url"]
+        if not isinstance(redis_url, str) or not redis_url.startswith("redis"):
+            return
+        self._pubsub = RedisPubSub(redis_url)
+        await self._pubsub.connect()
+
+        async def on_invalidate(event_payload: dict[str, Any]) -> None:
+            table = event_payload.get("table")
+            pks = event_payload.get("pks", [])
+            if not table:
+                return
+            transport = self._transport
+            assert transport is not None
+            if pks:
+                await invalidate_tags(transport, *(f"{table}:{pk}" for pk in pks))
+            else:
+                await _bump_table_version(transport, table)
+
+        await self._pubsub.listen(on_invalidate)
+        await self._pubsub.start()
+
+    async def _ensure_transport(self) -> None:
+        if self._transport is None:
+            self._transport = CashewsTransport.from_config(self._config)
+            await self._transport.connect()
+
+    async def _publish_invalidation(self, model: type[Any], pks: list[Any], action: str) -> None:
+        if self._pubsub is None:
+            return
+        await self._pubsub.publish(
+            {
+                "table": model.__tablename__,
+                "pks": pks,
+                "action": action,
+                "version": 1,
+            }
+        )
+
+    def _schedule_invalidation(self, model: type[Any], target: Any, action: str) -> None:
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        pk = extract_pk_from_instance(target)
+        task = asyncio.create_task(self.invalidate(model=model, pks=[pk]))
+        self._pending_tasks.add(task)
+        task.add_done_callback(self._pending_tasks.discard)
+
+    def _schedule_table_bump(self, model: type[Any]) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        transport = self._transport
+        if transport is None:
+            return
+        task = loop.create_task(_bump_table_version(transport, model))
+        self._pending_tasks.add(task)
+        task.add_done_callback(self._pending_tasks.discard)
+
+    def _handle_select(self, execute_state: Any) -> Any:
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return self._run_sync(resolve_cached_result(self, execute_state))
+        else:
+            return execute_state.invoke_statement()
+
+    def _handle_bulk_mutation(self, execute_state: Any) -> Any:
+        return self._run_sync(handle_bulk_mutation(self, execute_state))
+
+    @staticmethod
+    def _run_sync(coro: Any) -> Any:
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(coro)
+        raise RuntimeError("Sync wrappers cannot be used while an event loop is already running")
+
+    @classmethod
+    def _patch_async_get(cls) -> None:
+        if cls._async_get_patched:
+            return
+
+        cls._original_async_get = AsyncSession.get
+
+        async def patched_get(self: AsyncSession, entity: Any, ident: Any, **kwargs: Any) -> Any:
+            bind = self.sync_session.get_bind(mapper=entity)
+            manager = cls._engine_registry.get(id(bind))
+            if manager is None or not manager.is_enabled(entity, "get"):
+                return await cls._original_async_get(self, entity, ident, **kwargs)
+
+            unsupported = (
+                kwargs.get("options"),
+                kwargs.get("populate_existing"),
+                kwargs.get("with_for_update"),
+                kwargs.get("identity_token"),
+            )
+            if any(unsupported):
+                return await cls._original_async_get(self, entity, ident, **kwargs)
+
+            statement = manager._build_get_statement(
+                entity,
+                ident,
+                execution_options=kwargs.get("execution_options"),
+            )
+            result = await manager.execute(self, statement)
+            return result.scalar_one_or_none()
+
+        setattr(AsyncSession, "get", patched_get)  # noqa: B010
+        cls._async_get_patched = True
+
+    @classmethod
+    def _restore_async_get(cls) -> None:
+        if cls._async_get_patched and cls._original_async_get is not None:
+            setattr(AsyncSession, "get", cls._original_async_get)  # noqa: B010
+        cls._async_get_patched = False
+        cls._original_async_get = None
+
+    @staticmethod
+    def _build_get_statement(entity: type[Any], ident: Any, execution_options: Any = None) -> Any:
+        mapper = entity.__mapper__
+        pk_columns = list(mapper.primary_key)
+        values = [ident] if len(pk_columns) == 1 and not isinstance(ident, tuple) else list(ident)
+        statement = select(entity).where(*[column == value for column, value in zip(pk_columns, values, strict=False)])
+        if execution_options:
+            statement = statement.execution_options(**dict(execution_options))
+        return statement

sqlacache/pubsub/__init__.py

@@ -0,0 +1,3 @@
+"""Cross-process invalidation helpers."""
+
+__all__: list[str] = []

sqlacache/pubsub/redis.py

@@ -0,0 +1,89 @@
+"""Redis pub/sub adapter for invalidation events."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any
+
+from redis.asyncio import Redis
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+logger = logging.getLogger(__name__)
+
+
+class RedisPubSub:
+    """Listen for and publish invalidation events over Redis pub/sub."""
+
+    def __init__(self, redis_url: str, channel: str = "sqlacache:invalidate") -> None:
+        self._redis_url = redis_url
+        self._channel = channel
+        self._client: Redis | None = None
+        self._pubsub: Any | None = None
+        self._task: asyncio.Task[None] | None = None
+        self._callbacks: list[Callable[[dict[str, Any]], Awaitable[None]]] = []
+
+    async def connect(self) -> None:
+        self._client = Redis.from_url(self._redis_url)
+        self._pubsub = self._client.pubsub()
+        await self._pubsub.subscribe(self._channel)
+
+    async def listen(self, callback: Callable[[dict[str, Any]], Awaitable[None]]) -> None:
+        self._callbacks.append(callback)
+
+    async def start(self) -> None:
+        if self._task is None:
+            self._task = asyncio.create_task(self._listen_loop())
+
+    async def _listen_loop(self) -> None:
+        if self._pubsub is None:
+            return
+        try:
+            async for message in self._pubsub.listen():
+                if message.get("type") != "message":
+                    continue
+                raw_data = message.get("data")
+                try:
+                    if isinstance(raw_data, bytes):
+                        raw_data = raw_data.decode("utf-8")
+                    event = json.loads(raw_data)
+                except (TypeError, json.JSONDecodeError):
+                    logger.warning("Failed to parse pub/sub message: %r", message)
+                    continue
+
+                for callback in list(self._callbacks):
+                    try:
+                        await callback(event)
+                    except Exception:
+                        logger.exception("Error in pub/sub callback")
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception("Pub/sub listener crashed")
+
+    async def publish(self, event: dict[str, Any]) -> int:
+        if self._client is None:
+            raise RuntimeError("RedisPubSub is not connected")
+        message = json.dumps(event, default=str, sort_keys=True)
+        return int(await self._client.publish(self._channel, message))
+
+    async def stop(self) -> None:
+        if self._task is not None:
+            self._task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+        if self._pubsub is not None:
+            await self._pubsub.unsubscribe(self._channel)
+            await self._pubsub.aclose()
+            self._pubsub = None
+
+    async def disconnect(self) -> None:
+        await self.stop()
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None

sqlacache/serializers/__init__.py

@@ -0,0 +1,5 @@
+"""Serialization helpers."""
+
+from sqlacache.serializers.json import ModelJSONSerializer
+
+__all__ = ["ModelJSONSerializer"]

sqlacache/serializers/json.py

@@ -0,0 +1,24 @@
+"""Model-aware JSON serialization helpers."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from sqlalchemy import inspect
+
+
+class ModelJSONSerializer:
+    """Serialize ORM model instances using column attributes only."""
+
+    def serialize(self, obj: Any) -> str:
+        if isinstance(obj, list):
+            return json.dumps([self._instance_to_dict(item) for item in obj], default=str, sort_keys=True)
+        return json.dumps(self._instance_to_dict(obj), default=str, sort_keys=True)
+
+    def deserialize(self, payload: str) -> Any:
+        return json.loads(payload)
+
+    @staticmethod
+    def _instance_to_dict(instance: Any) -> dict[str, Any]:
+        return {attr.key: getattr(instance, attr.key) for attr in inspect(instance).mapper.column_attrs}

sqlacache/transport/__init__.py

@@ -0,0 +1,45 @@
+"""Transport abstractions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class CacheTransport(Protocol):
+    """Protocol for cache storage backends."""
+
+    async def connect(self) -> None:
+        """Initialize the transport."""
+
+    async def disconnect(self) -> None:
+        """Release backend resources."""
+
+    async def get(self, key: str) -> Any | None:
+        """Retrieve a cached value."""
+
+    async def set(
+        self,
+        key: str,
+        value: Any,
+        expire: int,
+        tags: Sequence[str] | None = None,
+    ) -> None:
+        """Store a cached value."""
+
+    async def delete(self, *keys: str) -> None:
+        """Delete one or more keys."""
+
+    async def delete_tags(self, *tags: str) -> None:
+        """Invalidate keys associated with tags."""
+
+    async def clear(self) -> None:
+        """Delete all cached entries."""
+
+    async def incr(self, key: str) -> int:
+        """Atomically increment an integer counter and return the new value."""
+
+    async def is_available(self) -> bool:
+        """Return whether the backend is usable."""

sqlacache/transport/cashews.py

@@ -0,0 +1,136 @@
+"""Cashews-based transport implementation."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from cashews import Cache
+
+from sqlacache.exceptions import ConfigError, TransportError
+
+logger = logging.getLogger(__name__)
+
+
+class CashewsTransport:
+    """Thin async wrapper around ``cashews.Cache``."""
+
+    def __init__(self, url: str, **kwargs: Any) -> None:
+        self._url = url
+        self._kwargs = dict(kwargs)
+        self._suppress = bool(self._kwargs.get("suppress", False))
+        self._cache = Cache()
+        self._connected = False
+
+    @classmethod
+    def from_config(cls, config: Mapping[str, Any]) -> CashewsTransport:
+        """Build a transport instance from normalized cache config."""
+
+        backend_config = config.get("backend", {})
+        if isinstance(backend_config, str):
+            url = backend_config
+            kwargs: dict[str, Any] = {}
+        elif isinstance(backend_config, Mapping):
+            url = backend_config.get("url", "mem://")
+            if not isinstance(url, str) or not url:
+                raise ConfigError("Backend mapping must define a non-empty 'url' string")
+            kwargs = {key: value for key, value in backend_config.items() if key != "url"}
+        else:
+            raise ConfigError("Backend must be configured as a URL string or mapping")
+
+        if "pickle_type" not in kwargs:
+            kwargs["pickle_type"] = config.get("serializer", "sqlalchemy")
+
+        return cls(url=url, **kwargs)
+
+    async def connect(self) -> None:
+        try:
+            self._cache.setup(self._url, **self._kwargs)
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache connect failed for %s: %s", self._url, exc)
+                return
+            raise TransportError(f"Cache connect failed for {self._url!r}") from exc
+        self._connected = True
+
+    async def disconnect(self) -> None:
+        try:
+            await self._cache.close()
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache disconnect failed for %s: %s", self._url, exc)
+                return
+            raise TransportError("Cache disconnect failed") from exc
+        finally:
+            self._connected = False
+
+    async def get(self, key: str) -> Any | None:
+        try:
+            return await self._cache.get(key)
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache get failed for key %s: %s", key, exc)
+                return None
+            raise TransportError(f"Cache get failed for key {key!r}") from exc
+
+    async def set(
+        self,
+        key: str,
+        value: Any,
+        expire: int,
+        tags: Sequence[str] | None = None,
+    ) -> None:
+        try:
+            await self._cache.set(key, value, expire=expire, tags=tags or ())
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache set failed for key %s: %s", key, exc)
+                return
+            raise TransportError(f"Cache set failed for key {key!r}") from exc
+
+    async def delete(self, *keys: str) -> None:
+        try:
+            for key in keys:
+                await self._cache.delete(key)
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache delete failed for keys %s: %s", keys, exc)
+                return
+            raise TransportError(f"Cache delete failed for keys {keys!r}") from exc
+
+    async def delete_tags(self, *tags: str) -> None:
+        try:
+            await self._cache.delete_tags(*tags)
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache delete_tags failed for tags %s: %s", tags, exc)
+                return
+            raise TransportError(f"Cache delete_tags failed for tags {tags!r}") from exc
+
+    async def clear(self) -> None:
+        try:
+            await self._cache.clear()
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache clear failed for %s: %s", self._url, exc)
+                return
+            raise TransportError("Cache clear failed") from exc
+
+    async def incr(self, key: str) -> int:
+        try:
+            return int(await self._cache.incr(key))
+        except Exception as exc:
+            if self._suppress:
+                logger.warning("Cache incr failed for key %s: %s", key, exc)
+                return 0
+            raise TransportError(f"Cache incr failed for key {key!r}") from exc
+
+    async def is_available(self) -> bool:
+        if not self._connected:
+            return False
+        try:
+            await self._cache.get("__sqlacache_healthcheck__")
+        except Exception:
+            return False
+        return True

sqlacache/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility helpers for sqlacache."""