hishel 0.1.5__py3-none-any.whl → 1.0.0b1__py3-none-any.whl

hishel/_core/_storages/_async_base.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import abc
+import time
+import typing as tp
+import uuid
+
+from ..models import Entry, Request, Response
+
+
+class AsyncBaseStorage(abc.ABC):
+    @abc.abstractmethod
+    async def create_entry(self, request: Request, response: Response, key: str, id_: uuid.UUID | None = None) -> Entry:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def get_entries(self, key: str) -> tp.List[Entry]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def update_entry(
+        self,
+        id: uuid.UUID,
+        new_entry: tp.Union[Entry, tp.Callable[[Entry], Entry]],
+    ) -> tp.Optional[Entry]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def remove_entry(self, id: uuid.UUID) -> None:
+        raise NotImplementedError()
+
+    async def close(self) -> None:
+        pass
+
+    def is_soft_deleted(self, pair: Entry) -> bool:
+        """
+        Check if a pair is soft deleted based on its metadata.
+
+        Args:
+            pair: The request pair to check.
+
+        Returns:
+            True if the pair is soft deleted, False otherwise.
+        """
+        return pair.meta.deleted_at is not None and pair.meta.deleted_at > 0
+
+    def is_safe_to_hard_delete(self, pair: Entry) -> bool:
+        """
+        Check if a pair is safe to hard delete based on its metadata.
+
+        If the pair has been soft deleted for more than 1 hour, it is considered safe to hard delete.
+
+        Args:
+            pair: The request pair to check.
+
+        Returns:
+            True if the pair is safe to hard delete, False otherwise.
+        """
+        return bool(pair.meta.deleted_at is not None and (pair.meta.deleted_at + 3600 < time.time()))
+
+    def mark_pair_as_deleted(self, pair: Entry) -> Entry:
+        """
+        Mark a pair as soft deleted by setting its deleted_at timestamp.
+
+        Args:
+            pair: The request pair to mark as deleted.
+        Returns:
+            The updated request pair with the deleted_at timestamp set.
+        """
+        pair.meta.deleted_at = time.time()
+        return pair

hishel/_core/_storages/_async_sqlite.py

@@ -0,0 +1,420 @@
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import replace
+from typing import (
+    Any,
+    AsyncIterable,
+    AsyncIterator,
+    Callable,
+    List,
+    Optional,
+    Union,
+)
+
+from hishel._core._storages._async_base import AsyncBaseStorage
+from hishel._core._storages._packing import pack, unpack
+from hishel._core.models import (
+    Entry,
+    EntryMeta,
+    Request,
+    Response,
+)
+from hishel._utils import ensure_cache_dict
+
+# Batch cleanup configuration
+# How often to run cleanup (seconds). Default: 1 hour.
+BATCH_CLEANUP_INTERVAL = 3600
+# How long to wait after storage creation before allowing the first cleanup (seconds)
+BATCH_CLEANUP_START_DELAY = 5 * 60
+# Number of rows to process per chunk when cleaning
+BATCH_CLEANUP_CHUNK_SIZE = 200
+
+
+try:
+    import anysqlite
+
+    class AsyncSqliteStorage(AsyncBaseStorage):
+        _COMPLETE_CHUNK_NUMBER = -1
+
+        def __init__(
+            self,
+            *,
+            connection: Optional[anysqlite.Connection] = None,
+            database_path: str = "hishel_cache.db",
+            default_ttl: Optional[float] = None,
+            refresh_ttl_on_access: bool = True,
+        ) -> None:
+            base_path = ensure_cache_dict()
+
+            self.connection = connection
+            self.database_path = base_path / database_path
+            self.default_ttl = default_ttl
+            self.refresh_ttl_on_access = refresh_ttl_on_access
+            self.last_cleanup = time.time() - BATCH_CLEANUP_INTERVAL + BATCH_CLEANUP_START_DELAY
+            # When this storage instance was created. Used to delay the first cleanup.
+            self._start_time = time.time()
+            self._initialized = False
+
+        async def _ensure_connection(self) -> anysqlite.Connection:
+            """Ensure connection is established and database is initialized."""
+            if self.connection is None:
+                self.connection = await anysqlite.connect(str(self.database_path))
+            if not self._initialized:
+                await self._initialize_database()
+                self._initialized = True
+            return self.connection
+
+        async def _initialize_database(self) -> None:
+            """Initialize the database schema."""
+            assert self.connection is not None
+            cursor = await self.connection.cursor()
+
+            # Table for storing request/response pairs
+            await cursor.execute("""
+                CREATE TABLE IF NOT EXISTS entries (
+                    id BLOB PRIMARY KEY,
+                    cache_key BLOB,
+                    data BLOB NOT NULL,
+                    created_at REAL NOT NULL,
+                    deleted_at REAL
+                )
+            """)
+
+            # Table for storing response stream chunks only
+            await cursor.execute("""
+                CREATE TABLE IF NOT EXISTS streams (
+                    entry_id BLOB NOT NULL,
+                    chunk_number INTEGER NOT NULL,
+                    chunk_data BLOB NOT NULL,
+                    PRIMARY KEY (entry_id, chunk_number),
+                    FOREIGN KEY (entry_id) REFERENCES entries(id) ON DELETE CASCADE
+                )
+            """)
+
+            # Indexes for performance
+            await cursor.execute("CREATE INDEX IF NOT EXISTS idx_entries_deleted_at ON entries(deleted_at)")
+            await cursor.execute("CREATE INDEX IF NOT EXISTS idx_entries_cache_key ON entries(cache_key)")
+
+            await self.connection.commit()
+
+        async def create_entry(
+            self, request: Request, response: Response, key: str, id_: uuid.UUID | None = None
+        ) -> Entry:
+            key_bytes = key.encode("utf-8")
+
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+
+            # Create a new entry directly with both request and response
+            pair_id = id_ if id_ is not None else uuid.uuid4()
+            pair_meta = EntryMeta(
+                created_at=time.time(),
+            )
+
+            assert isinstance(response.stream, (AsyncIterator, AsyncIterable))
+            response_with_stream = replace(
+                response,
+                stream=self._save_stream(response.stream, pair_id.bytes),
+            )
+
+            complete_entry = Entry(
+                id=pair_id,
+                request=request,
+                response=response_with_stream,
+                meta=pair_meta,
+                cache_key=key_bytes,
+            )
+
+            # Insert the complete entry into the database
+            await cursor.execute(
+                "INSERT INTO entries (id, cache_key, data, created_at, deleted_at) VALUES (?, ?, ?, ?, ?)",
+                (pair_id.bytes, key_bytes, pack(complete_entry, kind="pair"), pair_meta.created_at, None),
+            )
+            await connection.commit()
+
+            return complete_entry
+
+        async def get_entries(self, key: str) -> List[Entry]:
+            final_pairs: List[Entry] = []
+
+            now = time.time()
+            if now - self.last_cleanup >= BATCH_CLEANUP_INTERVAL:
+                try:
+                    await self._batch_cleanup()
+                except Exception:
+                    # don't let cleanup prevent reads; failures are non-fatal
+                    pass
+
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+            # Query entries directly by cache_key
+            await cursor.execute(
+                "SELECT id, data FROM entries WHERE cache_key = ?",
+                (key.encode("utf-8"),),
+            )
+
+            for row in await cursor.fetchall():
+                pair_data = unpack(row[1], kind="pair")
+
+                # Skip entries without a response (incomplete)
+                if not isinstance(pair_data, Entry) or pair_data.response is None:
+                    continue
+
+                final_pairs.append(pair_data)
+
+            pairs_with_streams: List[Entry] = []
+
+            # Only restore response streams from cache
+            for pair in final_pairs:
+                pairs_with_streams.append(
+                    replace(
+                        pair,
+                        response=replace(
+                            pair.response,
+                            stream=self._stream_data_from_cache(pair.id.bytes),
+                        ),
+                    )
+                )
+            return pairs_with_streams
+
+        async def update_entry(
+            self,
+            id: uuid.UUID,
+            new_pair: Union[Entry, Callable[[Entry], Entry]],
+        ) -> Optional[Entry]:
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+            await cursor.execute("SELECT data FROM entries WHERE id = ?", (id.bytes,))
+            result = await cursor.fetchone()
+
+            if result is None:
+                return None
+
+            pair = unpack(result[0], kind="pair")
+
+            # Skip entries without a response (incomplete)
+            if not isinstance(pair, Entry) or pair.response is None:
+                return None
+
+            if isinstance(new_pair, Entry):
+                complete_pair = new_pair
+            else:
+                complete_pair = new_pair(pair)
+
+            if pair.id != complete_pair.id:
+                raise ValueError("Pair ID mismatch")
+
+            await cursor.execute(
+                "UPDATE entries SET data = ? WHERE id = ?",
+                (pack(complete_pair, kind="pair"), id.bytes),
+            )
+
+            if pair.cache_key != complete_pair.cache_key:
+                await cursor.execute(
+                    "UPDATE entries SET cache_key = ? WHERE id = ?",
+                    (complete_pair.cache_key, complete_pair.id.bytes),
+                )
+
+            await connection.commit()
+
+            return complete_pair
+
+        async def remove_entry(self, id: uuid.UUID) -> None:
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+            await cursor.execute("SELECT data FROM entries WHERE id = ?", (id.bytes,))
+            result = await cursor.fetchone()
+
+            if result is None:
+                return None
+
+            pair = unpack(result[0], kind="pair")
+            await self._soft_delete_pair(pair, cursor)
+            await connection.commit()
+
+        async def _is_stream_complete(self, pair_id: uuid.UUID, cursor: anysqlite.Cursor) -> bool:
+            # Check if there's a completion marker (chunk_number = -1) for response stream
+            await cursor.execute(
+                "SELECT 1 FROM streams WHERE entry_id = ? AND chunk_number = ? LIMIT 1",
+                (pair_id.bytes, self._COMPLETE_CHUNK_NUMBER),
+            )
+            return await cursor.fetchone() is not None
+
+        async def _soft_delete_pair(
+            self,
+            pair: Entry,
+            cursor: anysqlite.Cursor,
+        ) -> None:
+            """
+            Mark the pair as deleted by setting the deleted_at timestamp.
+            """
+            marked_pair = self.mark_pair_as_deleted(pair)
+            await cursor.execute(
+                "UPDATE entries SET data = ?, deleted_at = ? WHERE id = ?",
+                (
+                    pack(marked_pair, kind="pair"),
+                    marked_pair.meta.deleted_at,
+                    pair.id.bytes,
+                ),
+            )
+
+        async def _is_pair_expired(self, pair: Entry, cursor: anysqlite.Cursor) -> bool:
+            """
+            Check if the pair is expired.
+            """
+            ttl = pair.request.metadata["hishel_ttl"] if "hishel_ttl" in pair.request.metadata else self.default_ttl
+            created_at = pair.meta.created_at
+            if ttl is None:
+                return False
+            return created_at + ttl < time.time()
+
+        async def _batch_cleanup(
+            self,
+        ) -> None:
+            """
+            Cleanup expired entries in the database.
+            """
+            should_mark_as_deleted: List[Entry] = []
+            should_hard_delete: List[Entry] = []
+
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+
+            # Process entries in chunks to avoid loading the entire table into memory.
+            chunk_size = BATCH_CLEANUP_CHUNK_SIZE
+            offset = 0
+            while True:
+                await cursor.execute(
+                    "SELECT id, data FROM entries LIMIT ? OFFSET ?",
+                    (chunk_size, offset),
+                )
+                rows = await cursor.fetchall()
+                if not rows:
+                    break
+
+                for row in rows:
+                    pair = unpack(row[1], kind="pair")
+                    if pair is None:
+                        continue
+
+                    # expired but not yet soft-deleted
+                    if await self._is_pair_expired(pair, cursor) and not self.is_soft_deleted(pair):
+                        should_mark_as_deleted.append(pair)
+
+                    # soft-deleted and safe to hard delete, or corrupted pair
+                    if (self.is_soft_deleted(pair) and self.is_safe_to_hard_delete(pair)) or await self._is_corrupted(
+                        pair, cursor
+                    ):
+                        should_hard_delete.append(pair)
+
+                # advance pagination
+                offset += len(rows)
+
+            for pair in should_mark_as_deleted:
+                await self._soft_delete_pair(pair, cursor)
+
+            for pair in should_hard_delete:
+                await self._hard_delete_pair(pair, cursor)
+
+            await connection.commit()
+
+        async def _is_corrupted(self, pair: Entry, cursor: anysqlite.Cursor) -> bool:
+            # if entry was created more than 1 hour ago and still has no response (incomplete)
+            if pair.meta.created_at + 3600 < time.time() and pair.response is None:
+                return True
+
+            # Check if response stream is complete for Entry with response
+            if (
+                isinstance(pair, Entry)
+                and pair.response is not None
+                and not await self._is_stream_complete(pair.id, cursor)
+            ):
+                return True
+            return False
+
+        async def _hard_delete_pair(self, pair: Entry, cursor: anysqlite.Cursor) -> None:
+            """
+            Permanently delete the pair from the database.
+            """
+            await cursor.execute("DELETE FROM entries WHERE id = ?", (pair.id.bytes,))
+
+            # Delete response stream for this entry
+            await self._delete_stream(pair.id.bytes, cursor)
+
+        async def _delete_stream(
+            self,
+            entry_id: bytes,
+            cursor: anysqlite.Cursor,
+        ) -> None:
+            """
+            Delete response stream associated with the given entry ID.
+            """
+            await cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
+
+        async def _save_stream(
+            self,
+            stream: AsyncIterator[bytes],
+            entry_id: bytes,
+        ) -> AsyncIterator[bytes]:
+            """
+            Wrapper around an async iterator that also saves the response data to the cache in chunks.
+            """
+            chunk_number = 0
+            content_length = 0
+            async for chunk in stream:
+                content_length += len(chunk)
+                connection = await self._ensure_connection()
+                cursor = await connection.cursor()
+                await cursor.execute(
+                    "INSERT INTO streams (entry_id, chunk_number, chunk_data) VALUES (?, ?, ?)",
+                    (entry_id, chunk_number, chunk),
+                )
+                await connection.commit()
+                chunk_number += 1
+                yield chunk
+
+            # Mark end of stream with chunk_number = -1
+            connection = await self._ensure_connection()
+            cursor = await connection.cursor()
+            await cursor.execute(
+                "INSERT INTO streams (entry_id, chunk_number, chunk_data) VALUES (?, ?, ?)",
+                (entry_id, self._COMPLETE_CHUNK_NUMBER, b""),
+            )
+            await connection.commit()
+
+        async def _stream_data_from_cache(
+            self,
+            entry_id: bytes,
+        ) -> AsyncIterator[bytes]:
+            """
+            Get an async iterator that yields the response stream data from the cache.
+            """
+            chunk_number = 0
+
+            connection = await self._ensure_connection()
+            while True:
+                cursor = await connection.cursor()
+                await cursor.execute(
+                    "SELECT chunk_data FROM streams WHERE entry_id = ? AND chunk_number = ?",
+                    (entry_id, chunk_number),
+                )
+                result = await cursor.fetchone()
+
+                if result is None:
+                    break
+                chunk = result[0]
+                # chunk_number = -1 is the completion marker with empty data
+                if chunk == b"":
+                    break
+                yield chunk
+                chunk_number += 1
+except ImportError:
+
+    class AsyncSqliteStorage:  # type: ignore[no-redef]
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            raise ImportError(
+                "The 'anysqlite' library is required to use the `AsyncSqliteStorage` integration. "
+                "Install hishel with 'pip install hishel[async]'."
+            )

hishel/_core/_storages/_packing.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Union, overload
+
+import msgpack
+from typing_extensions import Literal, cast
+
+from hishel._core._headers import Headers
+from hishel._core.models import EntryMeta, Request, Response
+
+
+def filter_out_hishel_metadata(data: Mapping[str, Any]) -> dict[str, Any]:
+    return {k: v for k, v in data.items() if not k.startswith("hishel_")}
+
+
+if TYPE_CHECKING:
+    from hishel import Entry
+
+
+@overload
+def pack(
+    value: "Entry",
+    /,
+    kind: Literal["pair"],
+) -> bytes: ...
+
+
+@overload
+def pack(
+    value: uuid.UUID,
+    /,
+    kind: Literal["entry_db_key_index"],
+) -> bytes: ...
+
+
+def pack(
+    value: Union["Entry", uuid.UUID],
+    /,
+    kind: Literal["pair", "entry_db_key_index"],
+) -> bytes:
+    from hishel import Entry
+
+    if kind == "entry_db_key_index":
+        assert isinstance(value, uuid.UUID)
+        return value.bytes
+    elif kind == "pair":
+        assert isinstance(value, Entry)
+        return cast(
+            bytes,
+            msgpack.packb(
+                {
+                    "id": value.id.bytes,
+                    "request": {
+                        "method": value.request.method,
+                        "url": value.request.url,
+                        "headers": value.request.headers._headers,
+                        "extra": filter_out_hishel_metadata(value.request.metadata),
+                    },
+                    "response": {
+                        "status_code": value.response.status_code,
+                        "headers": value.response.headers._headers,
+                        "extra": filter_out_hishel_metadata(value.response.metadata),
+                    },
+                    "meta": {
+                        "created_at": value.meta.created_at,
+                        "deleted_at": value.meta.deleted_at,
+                    },
+                    "cache_key": value.cache_key,
+                }
+            ),
+        )
+    assert False, f"Unexpected kind: {kind}"
+
+
+@overload
+def unpack(
+    value: bytes,
+    /,
+    kind: Literal["pair"],
+) -> "Entry": ...
+
+
+@overload
+def unpack(
+    value: bytes,
+    /,
+    kind: Literal["entry_db_key_index"],
+) -> uuid.UUID: ...
+
+
+@overload
+def unpack(
+    value: Optional[bytes],
+    /,
+    kind: Literal["pair"],
+) -> Optional["Entry"]: ...
+
+
+@overload
+def unpack(
+    value: Optional[bytes],
+    /,
+    kind: Literal["entry_db_key_index"],
+) -> Optional[uuid.UUID]: ...
+
+
+def unpack(
+    value: Optional[bytes],
+    /,
+    kind: Literal["pair", "entry_db_key_index"],
+) -> Union["Entry", uuid.UUID, None]:
+    from hishel import Entry
+
+    if value is None:
+        return None
+    if kind == "entry_db_key_index":
+        return uuid.UUID(bytes=value)
+    elif kind == "pair":
+        data = msgpack.unpackb(value)
+        id = uuid.UUID(bytes=data["id"])
+        return Entry(
+            id=id,
+            request=Request(
+                method=data["request"]["method"],
+                url=data["request"]["url"],
+                headers=Headers(data["request"]["headers"]),
+                metadata=data["request"]["extra"],
+                stream=iter([]),
+            ),
+            response=(
+                Response(
+                    status_code=data["response"]["status_code"],
+                    headers=Headers(data["response"]["headers"]),
+                    metadata=data["response"]["extra"],
+                    stream=iter([]),
+                )
+            ),
+            meta=EntryMeta(
+                created_at=data["meta"]["created_at"],
+                deleted_at=data["meta"]["deleted_at"],
+            ),
+            cache_key=data["cache_key"],
+        )

hishel/_core/_storages/_sync_base.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import abc
+import time
+import typing as tp
+import uuid
+
+from ..models import Entry, Request, Response
+
+
+class SyncBaseStorage(abc.ABC):
+    @abc.abstractmethod
+    def create_entry(self, request: Request, response: Response, key: str, id_: uuid.UUID | None = None) -> Entry:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_entries(self, key: str) -> tp.List[Entry]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def update_entry(
+        self,
+        id: uuid.UUID,
+        new_entry: tp.Union[Entry, tp.Callable[[Entry], Entry]],
+    ) -> tp.Optional[Entry]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def remove_entry(self, id: uuid.UUID) -> None:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        pass
+
+    def is_soft_deleted(self, pair: Entry) -> bool:
+        """
+        Check if a pair is soft deleted based on its metadata.
+
+        Args:
+            pair: The request pair to check.
+
+        Returns:
+            True if the pair is soft deleted, False otherwise.
+        """
+        return pair.meta.deleted_at is not None and pair.meta.deleted_at > 0
+
+    def is_safe_to_hard_delete(self, pair: Entry) -> bool:
+        """
+        Check if a pair is safe to hard delete based on its metadata.
+
+        If the pair has been soft deleted for more than 1 hour, it is considered safe to hard delete.
+
+        Args:
+            pair: The request pair to check.
+
+        Returns:
+            True if the pair is safe to hard delete, False otherwise.
+        """
+        return bool(pair.meta.deleted_at is not None and (pair.meta.deleted_at + 3600 < time.time()))
+
+    def mark_pair_as_deleted(self, pair: Entry) -> Entry:
+        """
+        Mark a pair as soft deleted by setting its deleted_at timestamp.
+
+        Args:
+            pair: The request pair to mark as deleted.
+        Returns:
+            The updated request pair with the deleted_at timestamp set.
+        """
+        pair.meta.deleted_at = time.time()
+        return pair