hishel 1.0.0.dev2__py3-none-any.whl → 1.0.0.dev3__py3-none-any.whl

hishel/_core/{_sync/_storages/_sqlite.py → _storages/_sync_sqlite.py}

@@ -9,21 +9,19 @@ from typing import (
     Iterator,
     Callable,
     List,
-    Literal,
     Optional,
     Union,
 )
 
-from hishel._core._base._storages._base import SyncBaseStorage, ensure_cache_dict
-from hishel._core._base._storages._packing import pack, unpack
+from hishel._core._storages._sync_base import SyncBaseStorage
+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 sqlite3
 
     class SyncSqliteStorage(SyncBaseStorage):
-        _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
             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
             cursor.execute("CREATE INDEX IF NOT EXISTS idx_entries_deleted_at ON entries(deleted_at)")
             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 = ?
 
             self.connection.commit()
 
-        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")
+        def create_entry(
+            self, request: Request, response: Response, key: str, id_: uuid.UUID | None = None
+        ) -> Entry:
+            key_bytes = key.encode("utf-8")
 
             connection = self._ensure_connection()
             cursor = connection.cursor()
-            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),
-            )
-            connection.commit()
-
-            assert isinstance(request.stream, Iterable), "Request stream must be an Iterable, 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)
-
-        def add_response(
-            self,
-            pair_id: uuid.UUID,
-            response: Response,
-            key: str | bytes,
-        ) -> CompletePair:
-            if isinstance(key, str):
-                key = key.encode("utf-8")
-
-            connection = self._ensure_connection()
-            cursor = connection.cursor()
-
-            # Get the existing pair
-            cursor.execute("SELECT data FROM entries WHERE id = ?", (pair_id.bytes,))
-            result = 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, (Iterator, Iterable))
-            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),
+            )
 
-            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
             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),
             )
             connection.commit()
 
-            return complete_pair
+            return complete_entry
 
-        def get_pairs(self, key: str) -> List[CompletePair]:
-            final_pairs: List[CompletePair] = []
+        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 = self._ensure_connection()
             cursor = connection.cursor()
             # Query entries directly by cache_key
-            cursor.execute("SELECT id, data FROM entries WHERE cache_key = ?", (key.encode("utf-8"),))
+            cursor.execute(
+                "SELECT id, data FROM entries WHERE cache_key = ?",
+                (key.encode("utf-8"),),
+            )
 
             for row in 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
 
-        def update_pair(
+        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 = self._ensure_connection()
             cursor = connection.cursor()
             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")
 
             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
 
-        def remove(self, id: uuid.UUID) -> None:
+        def remove_entry(self, id: uuid.UUID) -> None:
             connection = self._ensure_connection()
             cursor = connection.cursor()
             cursor.execute("SELECT data FROM entries WHERE id = ?", (id.bytes,))
@@ -272,28 +234,33 @@ try:
             self._soft_delete_pair(pair, cursor)
             connection.commit()
 
-        def _is_stream_complete(
-            self, kind: Literal["request", "response"], pair_id: uuid.UUID, cursor: sqlite3.Cursor
-        ) -> bool:
-            kind_id = self._STREAM_KIND[kind]
-            # Check if there's a completion marker (chunk_number = -1)
+        def _is_stream_complete(self, pair_id: uuid.UUID, cursor: sqlite3.Cursor) -> bool:
+            # Check if there's a completion marker (chunk_number = -1) for response stream
             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 cursor.fetchone() is not None
 
-        def _soft_delete_pair(self, pair: Union[CompletePair, IncompletePair], cursor: sqlite3.Cursor) -> None:
+        def _soft_delete_pair(
+            self,
+            pair: Entry,
+            cursor: sqlite3.Cursor,
+        ) -> None:
             """
             Mark the pair as deleted by setting the deleted_at timestamp.
             """
             marked_pair = self.mark_pair_as_deleted(pair)
             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,
+                ),
             )
 
-        def _is_pair_expired(self, pair: Pair, cursor: sqlite3.Cursor) -> bool:
+        def _is_pair_expired(self, pair: Entry, cursor: sqlite3.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 = self._ensure_connection()
             cursor = connection.cursor()
@@ -319,7 +286,10 @@ try:
             chunk_size = BATCH_CLEANUP_CHUNK_SIZE
             offset = 0
             while True:
-                cursor.execute("SELECT id, data FROM entries LIMIT ? OFFSET ?", (chunk_size, offset))
+                cursor.execute(
+                    "SELECT id, data FROM entries LIMIT ? OFFSET ?",
+                    (chunk_size, offset),
+                )
                 rows = cursor.fetchall()
                 if not rows:
                     break
@@ -350,61 +320,56 @@ try:
 
             connection.commit()
 
-        def _is_corrupted(self, pair: IncompletePair | CompletePair, cursor: sqlite3.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):
+        def _is_corrupted(self, pair: Entry, cursor: sqlite3.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 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 self._is_stream_complete(pair.id, cursor)
+            ):
                 return True
             return False
 
-        def _hard_delete_pair(self, pair: CompletePair | IncompletePair, cursor: sqlite3.Cursor) -> None:
+        def _hard_delete_pair(self, pair: Entry, cursor: sqlite3.Cursor) -> None:
             """
             Permanently delete the pair from the database.
             """
             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
             self._delete_stream(pair.id.bytes, cursor)
 
         def _delete_stream(
             self,
             entry_id: bytes,
             cursor: sqlite3.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":
-                cursor.execute(
-                    "DELETE FROM streams WHERE entry_id = ? AND kind = ?", (entry_id, self._STREAM_KIND["request"])
-                )
-            elif type == "response":
-                cursor.execute(
-                    "DELETE FROM streams WHERE entry_id = ? AND kind = ?", (entry_id, self._STREAM_KIND["response"])
-                )
-            elif type == "all":
-                cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
+            cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
 
         def _save_stream(
             self,
             stream: Iterator[bytes],
             entry_id: bytes,
-            kind: Literal["response", "request"],
         ) -> Iterator[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
             for chunk in stream:
+                content_length += len(chunk)
                 connection = self._ensure_connection()
                 cursor = connection.cursor()
                 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),
                 )
                 connection.commit()
                 chunk_number += 1
@@ -414,28 +379,26 @@ try:
             connection = self._ensure_connection()
             cursor = connection.cursor()
             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""),
             )
             connection.commit()
 
         def _stream_data_from_cache(
             self,
             entry_id: bytes,
-            kind: Literal["response", "request"],
         ) -> Iterator[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 = self._ensure_connection()
             while True:
                 cursor = connection.cursor()
                 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 = cursor.fetchone()
 
@@ -449,7 +412,7 @@ try:
                 chunk_number += 1
 except ImportError:
 
-    class SyncSqliteStorage(SyncBaseStorage):  # type: ignore[no-redef]
+    class SyncSqliteStorage:  # type: ignore[no-redef]
         def __init__(self, *args: Any, **kwargs: Any) -> None:
             raise ImportError(
                 "The 'sqlite3' library is required to use the `SyncSqliteStorage` integration. "

hishel/_core/models.py

@@ -64,7 +64,9 @@ class RequestMetadata(TypedDict, total=False):
     """
 
 
-def extract_metadata_from_headers(headers: Mapping[str, str]) -> RequestMetadata:
+def extract_metadata_from_headers(
+    headers: Mapping[str, str],
+) -> RequestMetadata:
     metadata: RequestMetadata = {}
     if "X-Hishel-Ttl" in headers:
         try:
@@ -146,34 +148,16 @@ class Response:
 
 
 @dataclass
-class PairMeta:
+class EntryMeta:
     created_at: float = field(default_factory=time.time)
     deleted_at: Optional[float] = None
 
 
 @dataclass
-class Pair:
+class Entry:
     id: uuid.UUID
     request: Request
-    meta: PairMeta
-
-
-# class used by storage
-@dataclass
-class IncompletePair(Pair):
-    extra: Mapping[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class CompletePair(Pair):
+    meta: EntryMeta
     response: Response
     cache_key: bytes
     extra: Mapping[str, Any] = field(default_factory=dict)
-
-    @classmethod
-    def create(
-        cls,
-        response: Response,
-        request: Request,
-    ) -> "CompletePair":  # pragma: nocover
-        return cls(id=uuid.uuid4(), request=request, response=response, meta=PairMeta(), cache_key=b"")

hishel/_sync_cache.py

@@ -4,7 +4,7 @@ import hashlib
 import logging
 import time
 from dataclasses import replace
-from typing import Iterator, Awaitable, Callable
+from typing import Iterable, Iterator, Awaitable, Callable
 
 from typing_extensions import assert_never
 
@@ -24,8 +24,9 @@ from hishel import (
     StoreAndUse,
     create_idle_state,
 )
-from hishel._core._spec import InvalidatePairs, vary_headers_match
-from hishel._core.models import CompletePair, ResponseMetadata
+from hishel._core._spec import InvalidateEntries, vary_headers_match
+from hishel._core.models import Entry, ResponseMetadata
+from hishel._utils import make_sync_iterator
 
 logger = logging.getLogger("hishel.integrations.clients")
 
@@ -42,12 +43,12 @@ class SyncCacheProxy:
 
     def __init__(
         self,
-        send_request: Callable[[Request], Response],
+        request_sender: Callable[[Request], Response],
         storage: SyncBaseStorage | None = None,
         cache_options: CacheOptions | None = None,
         ignore_specification: bool = False,
     ) -> None:
-        self.send_request = send_request
+        self.send_request = request_sender
         self.storage = storage if storage is not None else SyncSqliteStorage()
         self.cache_options = cache_options if cache_options is not None else CacheOptions()
         self.ignore_specification = ignore_specification
@@ -59,26 +60,30 @@ class SyncCacheProxy:
 
     def _get_key_for_request(self, request: Request) -> str:
         if request.metadata.get("hishel_body_key"):
-            assert isinstance(request.stream, Iterator)
+            assert isinstance(request.stream, (Iterator, Iterable))
             collected = b"".join([chunk for chunk in request.stream])
             hash_ = hashlib.sha256(collected).hexdigest()
+            request.stream = make_sync_iterator([collected])
             return f"{str(request.url)}-{hash_}"
-        return str(request.url)
+        return hashlib.sha256(str(request.url).encode("utf-8")).hexdigest()
 
-    def _maybe_refresh_pair_ttl(self, pair: CompletePair) -> None:
+    def _maybe_refresh_pair_ttl(self, pair: Entry) -> None:
         if pair.request.metadata.get("hishel_refresh_ttl_on_access"):
-            self.storage.update_pair(
+            self.storage.update_entry(
                 pair.id,
-                lambda complete_pair: replace(complete_pair, meta=replace(complete_pair.meta, created_at=time.time())),
+                lambda complete_pair: replace(
+                    complete_pair,
+                    meta=replace(complete_pair.meta, created_at=time.time()),
+                ),
             )
 
     def _handle_request_ignoring_spec(self, request: Request) -> Response:
         logger.debug("Trying to get cached response ignoring specification")
-        pairs = self.storage.get_pairs(self._get_key_for_request(request))
+        entries = self.storage.get_entries(self._get_key_for_request(request))
 
-        logger.debug(f"Found {len(pairs)} cached pairs for the request")
+        logger.debug(f"Found {len(entries)} cached entries for the request")
 
-        for pair in pairs:
+        for pair in entries:
             if (
                 str(pair.request.url) == str(request.url)
                 and pair.request.method == request.method
@@ -101,16 +106,15 @@ class SyncCacheProxy:
                 self._maybe_refresh_pair_ttl(pair)
                 return pair.response
 
-        incomplete_pair = self.storage.create_pair(
-            request,
-        )
-        response = self.send_request(incomplete_pair.request)
+        response = self.send_request(request)
 
         logger.debug("Storing response in cache ignoring specification")
-        complete_pair = self.storage.add_response(
-            incomplete_pair.id, response, self._get_key_for_request(request)
+        entry = self.storage.create_entry(
+            request,
+            response,
+            self._get_key_for_request(request),
         )
-        return complete_pair.response
+        return entry.response
 
     def _handle_request_respecting_spec(self, request: Request) -> Response:
         state: AnyState = create_idle_state("client", self.cache_options)
@@ -132,25 +136,26 @@ class SyncCacheProxy:
                 return state.pair.response
             elif isinstance(state, NeedToBeUpdated):
                 state = self._handle_update(state)
-            elif isinstance(state, InvalidatePairs):
-                state = self._handle_invalidate_pairs(state)
+            elif isinstance(state, InvalidateEntries):
+                state = self._handle_invalidate_entries(state)
             else:
                 assert_never(state)
 
         raise RuntimeError("Unreachable")
 
     def _handle_idle_state(self, state: IdleClient, request: Request) -> AnyState:
-        stored_pairs = self.storage.get_pairs(self._get_key_for_request(request))
-        return state.next(request, stored_pairs)
+        stored_entries = self.storage.get_entries(self._get_key_for_request(request))
+        return state.next(request, stored_entries)
 
     def _handle_cache_miss(self, state: CacheMiss) -> AnyState:
-        incomplete_pair = self.storage.create_pair(state.request)
-        response = self.send_request(incomplete_pair.request)
-        return state.next(response, incomplete_pair.id)
+        response = self.send_request(state.request)
+        return state.next(response)
 
     def _handle_store_and_use(self, state: StoreAndUse, request: Request) -> Response:
-        complete_pair = self.storage.add_response(
-            state.pair_id, state.response, self._get_key_for_request(request)
+        complete_pair = self.storage.create_entry(
+            request,
+            state.response,
+            self._get_key_for_request(request),
         )
         return complete_pair.response
 
@@ -159,16 +164,17 @@ class SyncCacheProxy:
         return state.next(revalidation_response)
 
     def _handle_update(self, state: NeedToBeUpdated) -> AnyState:
-        for pair in state.updating_pairs:
-            self.storage.update_pair(
-                pair.id,
+        for entry in state.updating_entries:
+            self.storage.update_entry(
+                entry.id,
                 lambda complete_pair: replace(
-                    complete_pair, response=replace(pair.response, headers=pair.response.headers)
+                    complete_pair,
+                    response=replace(entry.response, headers=entry.response.headers),
                 ),
             )
         return state.next()
 
-    def _handle_invalidate_pairs(self, state: InvalidatePairs) -> AnyState:
-        for pair_id in state.pair_ids:
-            self.storage.remove(pair_id)
+    def _handle_invalidate_entries(self, state: InvalidateEntries) -> AnyState:
+        for entry_id in state.entry_ids:
+            self.storage.remove_entry(entry_id)
         return state.next()