hishel 1.0.0.dev1__py3-none-any.whl → 1.0.0.dev3__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/{_async/_storages/_sqlite.py → _storages/_async_sqlite.py}

@@ -9,21 +9,19 @@ from typing import (
     AsyncIterator,
     Callable,
     List,
-    Literal,
     Optional,
     Union,
 )
 
-from hishel._core._base._storages._base import AsyncBaseStorage, ensure_cache_dict
-from hishel._core._base._storages._packing import pack, unpack
+from hishel._core._storages._async_base import AsyncBaseStorage
+from hishel._core._storages._packing import pack, unpack
 from hishel._core.models import (
-    CompletePair,
-    IncompletePair,
-    Pair,
-    PairMeta,
+    Entry,
+    EntryMeta,
     Request,
     Response,
 )
+from hishel._utils import ensure_cache_dict
 
 # Batch cleanup configuration
 # How often to run cleanup (seconds). Default: 1 hour.
@@ -38,7 +36,6 @@ try:
     import anysqlite
 
     class AsyncSqliteStorage(AsyncBaseStorage):
-        _STREAM_KIND = {"request": 0, "response": 1}
         _COMPLETE_CHUNK_NUMBER = -1
 
         def __init__(
@@ -85,14 +82,13 @@ try:
                 )
             """)
 
-            # Table for storing stream chunks
+            # Table for storing response stream chunks only
             await cursor.execute("""
                 CREATE TABLE IF NOT EXISTS streams (
                     entry_id BLOB NOT NULL,
-                    kind INTEGER NOT NULL,
                     chunk_number INTEGER NOT NULL,
                     chunk_data BLOB NOT NULL,
-                    PRIMARY KEY (entry_id, kind, chunk_number),
+                    PRIMARY KEY (entry_id, chunk_number),
                     FOREIGN KEY (entry_id) REFERENCES entries(id) ON DELETE CASCADE
                 )
             """)
@@ -100,85 +96,48 @@ try:
             # 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)")
-            # Note: PRIMARY KEY (entry_id, kind, chunk_number) already provides an index
-            # for queries like: entry_id = ? AND kind = ? AND chunk_number = ?
 
             await self.connection.commit()
 
-        async def create_pair(
-            self,
-            request: Request,
-            id: uuid.UUID | None = None,
-        ) -> IncompletePair:
-            pair_id = id if id is not None else uuid.uuid4()
-            pair_meta = PairMeta(
-                created_at=time.time(),
-            )
-
-            pair = IncompletePair(id=pair_id, request=request, meta=pair_meta)
-
-            packed_pair = pack(pair, kind="pair")
+        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()
-            await cursor.execute(
-                "INSERT INTO entries (id, cache_key, data, created_at, deleted_at) VALUES (?, ?, ?, ?, ?)",
-                (pair_id.bytes, None, packed_pair, pair_meta.created_at, None),
-            )
-            await connection.commit()
-
-            assert isinstance(request.stream, AsyncIterable), "Request stream must be an AsyncIterable, not Iterable"
 
-            request = Request(
-                method=request.method,
-                url=request.url,
-                headers=request.headers,
-                metadata=request.metadata,
-                stream=self._save_stream(request.stream, pair_id.bytes, "request"),
+            # 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(),
             )
 
-            return replace(pair, request=request)
-
-        async def add_response(
-            self,
-            pair_id: uuid.UUID,
-            response: Response,
-            key: str | bytes,
-        ) -> CompletePair:
-            if isinstance(key, str):
-                key = key.encode("utf-8")
-
-            connection = await self._ensure_connection()
-            cursor = await connection.cursor()
-
-            # Get the existing pair
-            await cursor.execute("SELECT data FROM entries WHERE id = ?", (pair_id.bytes,))
-            result = await cursor.fetchone()
-
-            if result is None:
-                raise ValueError(f"Entry with ID {pair_id} not found.")
-
-            pair = unpack(result[0], kind="pair")
-
             assert isinstance(response.stream, (AsyncIterator, AsyncIterable))
-            response = replace(response, stream=self._save_stream(response.stream, pair_id.bytes, "response"))
+            response_with_stream = replace(
+                response,
+                stream=self._save_stream(response.stream, pair_id.bytes),
+            )
 
-            await self._delete_stream(pair.id.bytes, cursor, type="response")
-            complete_pair = CompletePair(
-                id=pair.id, request=pair.request, response=response, meta=pair.meta, cache_key=key
+            complete_entry = Entry(
+                id=pair_id,
+                request=request,
+                response=response_with_stream,
+                meta=pair_meta,
+                cache_key=key_bytes,
             )
 
-            # Update the entry with the complete pair and set cache_key
+            # Insert the complete entry into the database
             await cursor.execute(
-                "UPDATE entries SET data = ?, cache_key = ? WHERE id = ?",
-                (pack(complete_pair, kind="pair"), key, pair_id.bytes),
+                "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_pair
+            return complete_entry
 
-        async def get_pairs(self, key: str) -> List[CompletePair]:
-            final_pairs: List[CompletePair] = []
+        async def get_entries(self, key: str) -> List[Entry]:
+            final_pairs: List[Entry] = []
 
             now = time.time()
             if now - self.last_cleanup >= BATCH_CLEANUP_INTERVAL:
@@ -191,39 +150,40 @@ try:
             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"),))
+            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")
 
-                if isinstance(pair_data, IncompletePair):
+                # 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[CompletePair] = []
+            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, "response"),
-                        ),
-                        request=replace(
-                            pair.request,
-                            stream=self._stream_data_from_cache(pair.id.bytes, "request"),
+                            stream=self._stream_data_from_cache(pair.id.bytes),
                         ),
                     )
                 )
             return pairs_with_streams
 
-        async def update_pair(
+        async def update_entry(
             self,
             id: uuid.UUID,
-            new_pair: Union[CompletePair, Callable[[CompletePair], CompletePair]],
-        ) -> Optional[CompletePair]:
+            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,))
@@ -234,10 +194,11 @@ try:
 
             pair = unpack(result[0], kind="pair")
 
-            if isinstance(pair, IncompletePair):
+            # Skip entries without a response (incomplete)
+            if not isinstance(pair, Entry) or pair.response is None:
                 return None
 
-            if isinstance(new_pair, CompletePair):
+            if isinstance(new_pair, Entry):
                 complete_pair = new_pair
             else:
                 complete_pair = new_pair(pair)
@@ -246,7 +207,8 @@ try:
                 raise ValueError("Pair ID mismatch")
 
             await cursor.execute(
-                "UPDATE entries SET data = ? WHERE id = ?", (pack(complete_pair, kind="pair"), id.bytes)
+                "UPDATE entries SET data = ? WHERE id = ?",
+                (pack(complete_pair, kind="pair"), id.bytes),
             )
 
             if pair.cache_key != complete_pair.cache_key:
@@ -259,7 +221,7 @@ try:
 
             return complete_pair
 
-        async def remove(self, id: uuid.UUID) -> None:
+        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,))
@@ -272,28 +234,33 @@ try:
             await self._soft_delete_pair(pair, cursor)
             await connection.commit()
 
-        async def _is_stream_complete(
-            self, kind: Literal["request", "response"], pair_id: uuid.UUID, cursor: anysqlite.Cursor
-        ) -> bool:
-            kind_id = self._STREAM_KIND[kind]
-            # Check if there's a completion marker (chunk_number = -1)
+        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 kind = ? AND chunk_number = ? LIMIT 1",
-                (pair_id.bytes, kind_id, self._COMPLETE_CHUNK_NUMBER),
+                "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: Union[CompletePair, IncompletePair], cursor: anysqlite.Cursor) -> 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),
+                (
+                    pack(marked_pair, kind="pair"),
+                    marked_pair.meta.deleted_at,
+                    pair.id.bytes,
+                ),
             )
 
-        async def _is_pair_expired(self, pair: Pair, cursor: anysqlite.Cursor) -> bool:
+        async def _is_pair_expired(self, pair: Entry, cursor: anysqlite.Cursor) -> bool:
             """
             Check if the pair is expired.
             """
@@ -307,10 +274,10 @@ try:
             self,
         ) -> None:
             """
-            Cleanup expired pairs in the database.
+            Cleanup expired entries in the database.
             """
-            should_mark_as_deleted: List[Union[CompletePair, IncompletePair]] = []
-            should_hard_delete: List[Union[CompletePair, IncompletePair]] = []
+            should_mark_as_deleted: List[Entry] = []
+            should_hard_delete: List[Entry] = []
 
             connection = await self._ensure_connection()
             cursor = await connection.cursor()
@@ -319,7 +286,10 @@ try:
             chunk_size = BATCH_CLEANUP_CHUNK_SIZE
             offset = 0
             while True:
-                await cursor.execute("SELECT id, data FROM entries LIMIT ? OFFSET ?", (chunk_size, offset))
+                await cursor.execute(
+                    "SELECT id, data FROM entries LIMIT ? OFFSET ?",
+                    (chunk_size, offset),
+                )
                 rows = await cursor.fetchall()
                 if not rows:
                     break
@@ -350,61 +320,56 @@ try:
 
             await connection.commit()
 
-        async def _is_corrupted(self, pair: IncompletePair | CompletePair, cursor: anysqlite.Cursor) -> bool:
-            # if pair was created more than 1 hour ago and still not completed
-            if pair.meta.created_at + 3600 < time.time() and isinstance(pair, IncompletePair):
+        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
 
-            if isinstance(pair, CompletePair) and not await self._is_stream_complete("request", pair.id, cursor):
+            # 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: CompletePair | IncompletePair, cursor: anysqlite.Cursor) -> None:
+        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 all streams (both request and response) for this entry
+            # 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,
-            type: Literal["request", "response", "all"] = "all",
         ) -> None:
             """
-            Delete all streams (both request and response) associated with the given entry ID.
+            Delete response stream associated with the given entry ID.
             """
-            if type == "request":
-                await cursor.execute(
-                    "DELETE FROM streams WHERE entry_id = ? AND kind = ?", (entry_id, self._STREAM_KIND["request"])
-                )
-            elif type == "response":
-                await cursor.execute(
-                    "DELETE FROM streams WHERE entry_id = ? AND kind = ?", (entry_id, self._STREAM_KIND["response"])
-                )
-            elif type == "all":
-                await cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
+            await cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
 
         async def _save_stream(
             self,
             stream: AsyncIterator[bytes],
             entry_id: bytes,
-            kind: Literal["response", "request"],
         ) -> AsyncIterator[bytes]:
             """
-            Wrapper around an async iterator that also saves the data to the cache in chunks.
+            Wrapper around an async iterator that also saves the response data to the cache in chunks.
             """
-            kind_id = self._STREAM_KIND[kind]
             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, kind, chunk_number, chunk_data) VALUES (?, ?, ?, ?)",
-                    (entry_id, kind_id, chunk_number, chunk),
+                    "INSERT INTO streams (entry_id, chunk_number, chunk_data) VALUES (?, ?, ?)",
+                    (entry_id, chunk_number, chunk),
                 )
                 await connection.commit()
                 chunk_number += 1
@@ -414,28 +379,26 @@ try:
             connection = await self._ensure_connection()
             cursor = await connection.cursor()
             await cursor.execute(
-                "INSERT INTO streams (entry_id, kind, chunk_number, chunk_data) VALUES (?, ?, ?, ?)",
-                (entry_id, kind_id, self._COMPLETE_CHUNK_NUMBER, b""),
+                "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,
-            kind: Literal["response", "request"],
         ) -> AsyncIterator[bytes]:
             """
-            Get an async iterator that yields the stream data from the cache.
+            Get an async iterator that yields the response stream data from the cache.
             """
-            kind_id = self._STREAM_KIND[kind]
             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 kind = ? AND chunk_number = ?",
-                    (entry_id, kind_id, chunk_number),
+                    "SELECT chunk_data FROM streams WHERE entry_id = ? AND chunk_number = ?",
+                    (entry_id, chunk_number),
                 )
                 result = await cursor.fetchone()
 
@@ -449,7 +412,7 @@ try:
                 chunk_number += 1
 except ImportError:
 
-    class AsyncSqliteStorage(AsyncBaseStorage):  # type: ignore[no-redef]
+    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. "

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"],
+        )