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

hishel/_core/_storages/_sync_sqlite.py

@@ -0,0 +1,420 @@
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import replace
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    Callable,
+    List,
+    Optional,
+    Union,
+)
+
+from hishel._core._storages._sync_base import SyncBaseStorage
+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 sqlite3
+
+    class SyncSqliteStorage(SyncBaseStorage):
+        _COMPLETE_CHUNK_NUMBER = -1
+
+        def __init__(
+            self,
+            *,
+            connection: Optional[sqlite3.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
+
+        def _ensure_connection(self) -> sqlite3.Connection:
+            """Ensure connection is established and database is initialized."""
+            if self.connection is None:
+                self.connection = sqlite3.connect(str(self.database_path))
+            if not self._initialized:
+                self._initialize_database()
+                self._initialized = True
+            return self.connection
+
+        def _initialize_database(self) -> None:
+            """Initialize the database schema."""
+            assert self.connection is not None
+            cursor = self.connection.cursor()
+
+            # Table for storing request/response pairs
+            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
+            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
+            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)")
+
+            self.connection.commit()
+
+        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()
+
+            # 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, (Iterator, Iterable))
+            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
+            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),
+            )
+            connection.commit()
+
+            return complete_entry
+
+        def get_entries(self, key: str) -> List[Entry]:
+            final_pairs: List[Entry] = []
+
+            now = time.time()
+            if now - self.last_cleanup >= BATCH_CLEANUP_INTERVAL:
+                try:
+                    self._batch_cleanup()
+                except Exception:
+                    # don't let cleanup prevent reads; failures are non-fatal
+                    pass
+
+            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"),),
+            )
+
+            for row in 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
+
+        def update_entry(
+            self,
+            id: uuid.UUID,
+            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,))
+            result = 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")
+
+            cursor.execute(
+                "UPDATE entries SET data = ? WHERE id = ?",
+                (pack(complete_pair, kind="pair"), id.bytes),
+            )
+
+            if pair.cache_key != complete_pair.cache_key:
+                cursor.execute(
+                    "UPDATE entries SET cache_key = ? WHERE id = ?",
+                    (complete_pair.cache_key, complete_pair.id.bytes),
+                )
+
+            connection.commit()
+
+            return complete_pair
+
+        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,))
+            result = cursor.fetchone()
+
+            if result is None:
+                return None
+
+            pair = unpack(result[0], kind="pair")
+            self._soft_delete_pair(pair, cursor)
+            connection.commit()
+
+        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 chunk_number = ? LIMIT 1",
+                (pair_id.bytes, self._COMPLETE_CHUNK_NUMBER),
+            )
+            return cursor.fetchone() is not 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,
+                ),
+            )
+
+        def _is_pair_expired(self, pair: Entry, cursor: sqlite3.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()
+
+        def _batch_cleanup(
+            self,
+        ) -> None:
+            """
+            Cleanup expired entries in the database.
+            """
+            should_mark_as_deleted: List[Entry] = []
+            should_hard_delete: List[Entry] = []
+
+            connection = self._ensure_connection()
+            cursor = connection.cursor()
+
+            # Process entries in chunks to avoid loading the entire table into memory.
+            chunk_size = BATCH_CLEANUP_CHUNK_SIZE
+            offset = 0
+            while True:
+                cursor.execute(
+                    "SELECT id, data FROM entries LIMIT ? OFFSET ?",
+                    (chunk_size, offset),
+                )
+                rows = 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 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 self._is_corrupted(
+                        pair, cursor
+                    ):
+                        should_hard_delete.append(pair)
+
+                # advance pagination
+                offset += len(rows)
+
+            for pair in should_mark_as_deleted:
+                self._soft_delete_pair(pair, cursor)
+
+            for pair in should_hard_delete:
+                self._hard_delete_pair(pair, cursor)
+
+            connection.commit()
+
+        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
+
+            # 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: Entry, cursor: sqlite3.Cursor) -> None:
+            """
+            Permanently delete the pair from the database.
+            """
+            cursor.execute("DELETE FROM entries WHERE id = ?", (pair.id.bytes,))
+
+            # Delete response stream for this entry
+            self._delete_stream(pair.id.bytes, cursor)
+
+        def _delete_stream(
+            self,
+            entry_id: bytes,
+            cursor: sqlite3.Cursor,
+        ) -> None:
+            """
+            Delete response stream associated with the given entry ID.
+            """
+            cursor.execute("DELETE FROM streams WHERE entry_id = ?", (entry_id,))
+
+        def _save_stream(
+            self,
+            stream: Iterator[bytes],
+            entry_id: bytes,
+        ) -> Iterator[bytes]:
+            """
+            Wrapper around an async iterator that also saves the response data to the cache in chunks.
+            """
+            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, chunk_number, chunk_data) VALUES (?, ?, ?)",
+                    (entry_id, chunk_number, chunk),
+                )
+                connection.commit()
+                chunk_number += 1
+                yield chunk
+
+            # Mark end of stream with chunk_number = -1
+            connection = self._ensure_connection()
+            cursor = connection.cursor()
+            cursor.execute(
+                "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,
+        ) -> Iterator[bytes]:
+            """
+            Get an async iterator that yields the response stream data from the cache.
+            """
+            chunk_number = 0
+
+            connection = self._ensure_connection()
+            while True:
+                cursor = connection.cursor()
+                cursor.execute(
+                    "SELECT chunk_data FROM streams WHERE entry_id = ? AND chunk_number = ?",
+                    (entry_id, chunk_number),
+                )
+                result = 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 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. "
+                "Install hishel with 'pip install hishel[async]'."
+            )

hishel/{beta/_core → _core}/models.py

@@ -5,14 +5,18 @@ import uuid
 from dataclasses import dataclass, field
 from typing import (
     Any,
+    AsyncIterable,
     AsyncIterator,
+    Iterable,
     Iterator,
     Mapping,
     Optional,
     TypedDict,
+    cast,
 )
 
-from hishel.beta._core._headers import Headers
+from hishel._core._headers import Headers
+from hishel._utils import make_async_iterator, make_sync_iterator
 
 
 class AnyIterable:
@@ -64,7 +68,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:
@@ -94,33 +100,71 @@ class Request:
     stream: Iterator[bytes] | AsyncIterator[bytes] = field(default_factory=lambda: iter(AnyIterable()))
     metadata: RequestMetadata | Mapping[str, Any] = field(default_factory=dict)
 
-    def iter_stream(self) -> Iterator[bytes]:
-        if isinstance(self.stream, Iterator):
-            return self.stream
+    def _iter_stream(self) -> Iterator[bytes]:
+        if hasattr(self, "collected_body"):
+            yield getattr(self, "collected_body")
+            return
+        if isinstance(self.stream, (Iterator, Iterable)):
+            yield from self.stream
+            return
         raise TypeError("Request stream is not an Iterator")
 
-    async def aiter_stream(self) -> AsyncIterator[bytes]:
-        if isinstance(self.stream, AsyncIterator):
+    async def _aiter_stream(self) -> AsyncIterator[bytes]:
+        if hasattr(self, "collected_body"):
+            yield getattr(self, "collected_body")
+            return
+        if isinstance(self.stream, (AsyncIterator, AsyncIterable)):
             async for chunk in self.stream:
                 yield chunk
+            return
         else:
             raise TypeError("Request stream is not an AsyncIterator")
 
+    def read(self) -> bytes:
+        """
+        Synchronously reads the entire request body without consuming the stream.
+        """
+        if not isinstance(self.stream, Iterator):
+            raise TypeError("Request stream is not an Iterator")
+
+        if hasattr(self, "collected_body"):
+            return cast(bytes, getattr(self, "collected_body"))
+
+        collected = b"".join([chunk for chunk in self.stream])
+        setattr(self, "collected_body", collected)
+        self.stream = make_sync_iterator([collected])
+        return collected
+
+    async def aread(self) -> bytes:
+        """
+        Asynchronously reads the entire request body without consuming the stream.
+        """
+        if not isinstance(self.stream, AsyncIterator):
+            raise TypeError("Request stream is not an AsyncIterator")
+
+        if hasattr(self, "collected_body"):
+            return cast(bytes, getattr(self, "collected_body"))
+
+        collected = b"".join([chunk async for chunk in self.stream])
+        setattr(self, "collected_body", collected)
+        self.stream = make_async_iterator([collected])
+        return collected
+
 
 class ResponseMetadata(TypedDict, total=False):
     # All the names here should be prefixed with "hishel_" to avoid collisions with user data
-    hishel_from_cache: bool | None
+    hishel_from_cache: bool
     """Indicates whether the response was served from cache."""
 
-    hishel_revalidated: bool | None
+    hishel_revalidated: bool
     """Indicates whether the response was revalidated with the origin server."""
 
-    hishel_spec_ignored: bool | None
-    """Indicates whether the caching specification was ignored for this response."""
-
-    hishel_stored: bool | None
+    hishel_stored: bool
     """Indicates whether the response was stored in cache."""
 
+    hishel_created_at: float
+    """Timestamp when the response was cached."""
+
 
 @dataclass
 class Response:
@@ -129,48 +173,67 @@ class Response:
     stream: Iterator[bytes] | AsyncIterator[bytes] = field(default_factory=lambda: iter(AnyIterable()))
     metadata: ResponseMetadata | Mapping[str, Any] = field(default_factory=dict)
 
-    def iter_stream(self) -> Iterator[bytes]:
+    def _iter_stream(self) -> Iterator[bytes]:
+        if hasattr(self, "collected_body"):
+            yield getattr(self, "collected_body")
+            return
         if isinstance(self.stream, Iterator):
-            return self.stream
+            yield from self.stream
+            return
         raise TypeError("Response stream is not an Iterator")
 
-    async def aiter_stream(self) -> AsyncIterator[bytes]:
+    async def _aiter_stream(self) -> AsyncIterator[bytes]:
+        if hasattr(self, "collected_body"):
+            yield getattr(self, "collected_body")
+            return
         if isinstance(self.stream, AsyncIterator):
             async for chunk in self.stream:
                 yield chunk
         else:
             raise TypeError("Response stream is not an AsyncIterator")
 
+    def read(self) -> bytes:
+        """
+        Synchronously reads the entire request body without consuming the stream.
+        """
+        if not isinstance(self.stream, Iterator):
+            raise TypeError("Request stream is not an Iterator")
+
+        if hasattr(self, "collected_body"):
+            return cast(bytes, getattr(self, "collected_body"))
+
+        collected = b"".join([chunk for chunk in self.stream])
+        setattr(self, "collected_body", collected)
+        self.stream = make_sync_iterator([collected])
+        return collected
+
+    async def aread(self) -> bytes:
+        """
+        Asynchronously reads the entire request body without consuming the stream.
+        """
+        if not isinstance(self.stream, AsyncIterator):
+            raise TypeError("Request stream is not an AsyncIterator")
+
+        if hasattr(self, "collected_body"):
+            return cast(bytes, getattr(self, "collected_body"))
+
+        collected = b"".join([chunk async for chunk in self.stream])
+        setattr(self, "collected_body", collected)
+        self.stream = make_async_iterator([collected])
+        return collected
+
 
 @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/_policies.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import abc
+import typing as t
+from dataclasses import dataclass, field
+from typing import Generic
+
+from hishel import Request, Response
+from hishel._core._spec import (
+    CacheOptions,
+)
+
+logger = __import__("logging").getLogger(__name__)
+
+T = t.TypeVar("T", Request, Response)
+
+
+class CachePolicy(abc.ABC):
+    use_body_key: bool = False
+    """Whether to include request body in cache key calculation."""
+
+
+class BaseFilter(abc.ABC, Generic[T]):
+    @abc.abstractmethod
+    def needs_body(self) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def apply(self, item: T, body: bytes | None) -> bool:
+        pass
+
+
+@dataclass
+class SpecificationPolicy(CachePolicy):
+    """
+    Caching policy that respects HTTP caching specification.
+    """
+
+    cache_options: CacheOptions = field(default_factory=CacheOptions)
+
+
+@dataclass
+class FilterPolicy(CachePolicy):
+    """
+    Caching policy that applies user-defined filtering logic.
+    """
+
+    request_filters: list[BaseFilter[Request]] = field(default_factory=list)
+    response_filters: list[BaseFilter[Response]] = field(default_factory=list)