kachedb 0.1.0a2__py3-none-any.whl

kachedb/__init__.py

@@ -0,0 +1,71 @@
+"""
+KacheDB — Python client for the KacheDB zero-copy storage engine.
+
+Install::
+
+    pip install kachedb
+    pip install kachedb[torch]  # for PyTorch tensor support
+
+Quickstart::
+
+    from kachedb import KacheClient
+
+    with KacheClient(host="127.0.0.1", port=6379) as client:
+        client.set("user:1", "alice", ex=3600)
+        print(client.get("user:1"))  # b"alice"
+
+Async::
+
+    from kachedb import AsyncKacheClient
+
+    async with AsyncKacheClient() as client:
+        await client.set("key", "value")
+        print(await client.get("key"))
+"""
+
+from ._version import __version__
+from .async_client import AsyncKacheClient
+from .client import KacheClient
+from .connection import Connection
+from .descriptor import TENSOR_DESCRIPTOR_MAGIC, TensorBlockDescriptor, TensorDType
+from .exceptions import (
+    ConnectionError,
+    KacheDBError,
+    PoolExhaustedError,
+    ProtocolError,
+    ResponseError,
+    TimeoutError,
+)
+from .pipeline import AsyncPipeline, Pipeline
+from .pool import AsyncConnectionPool, ConnectionPool
+from .tensor import attach_shm, detach_all, read_tensor, read_torch_tensor
+
+__all__ = [
+    "TENSOR_DESCRIPTOR_MAGIC",
+    "AsyncConnectionPool",
+    "AsyncKacheClient",
+    "AsyncPipeline",
+    # Connection
+    "Connection",
+    "ConnectionError",
+    "ConnectionPool",
+    # Clients
+    "KacheClient",
+    # Exceptions
+    "KacheDBError",
+    # Pipeline
+    "Pipeline",
+    "PoolExhaustedError",
+    "ProtocolError",
+    "ResponseError",
+    # Tensor / Zero-Copy
+    "TensorBlockDescriptor",
+    "TensorDType",
+    "TimeoutError",
+    # Version
+    "__version__",
+    "attach_shm",
+    "detach_all",
+    "read_tensor",
+    "read_torch_tensor",
+]

kachedb/_version.py

@@ -0,0 +1,3 @@
+"""Single-source version for kachedb Python SDK."""
+
+__version__ = "0.1.0a2"

kachedb/async_client.py

@@ -0,0 +1,169 @@
+"""
+Async KacheDB client using ``asyncio``.
+
+Provides the same high-level API as :class:`~kachedb.client.KacheClient`
+but uses non-blocking ``asyncio`` streams for use in async ML inference
+pipelines (vLLM, SGLang, etc.).
+
+Usage::
+
+    import asyncio
+    from kachedb import AsyncKacheClient
+
+    async def main():
+        async with AsyncKacheClient() as client:
+            await client.set("user:1", "alice", ex=3600)
+            print(await client.get("user:1"))
+
+    asyncio.run(main())
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .exceptions import ConnectionError
+from .pipeline import AsyncPipeline
+from .pool import AsyncConnectionPool
+from .resp import AsyncRespReader, RespValue, encode_command
+
+if TYPE_CHECKING:
+    import asyncio
+
+
+class AsyncKacheClient:
+    """High-level async KacheDB client.
+
+    Parameters
+    ----------
+    host : str
+        Server hostname or IP address.
+    port : int
+        Server TCP port.
+    decode_responses : bool
+        If ``True``, decode byte responses to UTF-8 strings.
+    max_connections : int
+        Maximum async connection pool size.
+    """
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 6379,
+        *,
+        decode_responses: bool = False,
+        max_connections: int = 10,
+    ) -> None:
+        self.host = host
+        self.port = port
+        self.decode_responses = decode_responses
+        self._pool = AsyncConnectionPool(
+            host=host,
+            port=port,
+            max_connections=max_connections,
+            decode_responses=decode_responses,
+        )
+        self._reader: asyncio.StreamReader | None = None
+        self._writer: asyncio.StreamWriter | None = None
+        self._resp_reader: AsyncRespReader | None = None
+
+    async def connect(self) -> AsyncKacheClient:
+        """Establish a dedicated async connection."""
+        self._reader, self._writer, self._resp_reader = await self._pool.get_connection()
+        return self
+
+    async def close(self) -> None:
+        """Release the connection back to the pool."""
+        if self._writer is not None and self._reader is not None and self._resp_reader is not None:
+            await self._pool.release_connection(self._reader, self._writer, self._resp_reader)
+            self._reader = None
+            self._writer = None
+            self._resp_reader = None
+
+    async def disconnect_all(self) -> None:
+        """Close all pooled connections."""
+        await self._pool.disconnect_all()
+
+    async def __aenter__(self) -> AsyncKacheClient:
+        return await self.connect()
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self.close()
+
+    # ── Internal Helpers ──────────────────────────────────────────────────
+
+    async def _execute(self, *args: str | bytes) -> RespValue:
+        """Execute a single command and return the response."""
+        if self._writer is None or self._resp_reader is None:
+            raise ConnectionError("Not connected to KacheDB")
+
+        data = encode_command(list(args))
+        self._writer.write(data)
+        await self._writer.drain()
+
+        response = await self._resp_reader.read_response()
+
+        if self.decode_responses and isinstance(response, bytes):
+            return response.decode("utf-8")
+
+        return response
+
+    # ── Redis-Compatible Commands ─────────────────────────────────────────
+
+    async def ping(self, message: str | None = None) -> str:
+        """Send ``PING`` and return ``PONG`` or the echoed message."""
+        args: list[str | bytes] = ["PING"]
+        if message is not None:
+            args.append(message)
+        result = await self._execute(*args)
+        return str(result) if result is not None else "PONG"
+
+    async def get(self, key: str | bytes) -> bytes | str | None:
+        """Retrieve the value for *key*."""
+        return await self._execute("GET", key)  # type: ignore[return-value]
+
+    async def set(
+        self,
+        key: str | bytes,
+        value: str | bytes,
+        *,
+        ex: int | None = None,
+        px: int | None = None,
+    ) -> bool:
+        """Store *value* under *key* with an optional TTL."""
+        args: list[str | bytes] = ["SET", key, value]
+        if ex is not None:
+            args.extend(["EX", str(ex)])
+        elif px is not None:
+            args.extend(["PX", str(px)])
+        result = await self._execute(*args)
+        return result == "OK"
+
+    async def mget(self, *keys: str | bytes) -> list[bytes | str | None]:
+        """Batch-retrieve values for multiple keys."""
+        if not keys:
+            return []
+        result = await self._execute("MGET", *keys)
+        return result if isinstance(result, list) else []  # type: ignore[return-value]
+
+    async def delete(self, *keys: str | bytes) -> int:
+        """Delete one or more keys."""
+        if not keys:
+            return 0
+        result = await self._execute("DEL", *keys)
+        return int(result) if isinstance(result, int) else 0
+
+    async def exists(self, *keys: str | bytes) -> int:
+        """Check existence of one or more keys."""
+        if not keys:
+            return 0
+        result = await self._execute("EXISTS", *keys)
+        return int(result) if isinstance(result, int) else 0
+
+    # ── Pipeline ──────────────────────────────────────────────────────────
+
+    def pipeline(self) -> AsyncPipeline:
+        """Create an async pipeline for batching commands."""
+        if self._writer is None or self._resp_reader is None:
+            raise ConnectionError("Not connected to KacheDB")
+        return AsyncPipeline(self._writer, self._resp_reader)

kachedb/client.py

@@ -0,0 +1,216 @@
+"""
+Synchronous KacheDB client.
+
+Provides a high-level, Redis-like API for communicating with a KacheDB
+daemon over TCP using the RESP2 wire protocol.
+
+Usage::
+
+    from kachedb import KacheClient
+
+    with KacheClient() as client:
+        client.set("user:1", "alice", ex=3600)
+        print(client.get("user:1"))  # b"alice"
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .pipeline import Pipeline
+from .pool import ConnectionPool
+
+if TYPE_CHECKING:
+    from .connection import Connection
+    from .resp import RespValue
+
+
+class KacheClient:
+    """High-level synchronous KacheDB client.
+
+    Supports both direct connection mode and connection pool mode.
+
+    Parameters
+    ----------
+    host : str
+        Server hostname or IP address.
+    port : int
+        Server TCP port.
+    socket_timeout : float | None
+        Socket timeout in seconds.
+    decode_responses : bool
+        If ``True``, decode byte responses to UTF-8 strings.
+    max_connections : int
+        Maximum pool size.  Set to ``1`` to disable pooling.
+    """
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 6379,
+        *,
+        socket_timeout: float | None = 5.0,
+        decode_responses: bool = False,
+        max_connections: int = 10,
+    ) -> None:
+        self.host = host
+        self.port = port
+        self._pool = ConnectionPool(
+            host=host,
+            port=port,
+            max_connections=max_connections,
+            socket_timeout=socket_timeout,
+            decode_responses=decode_responses,
+        )
+        self._conn: Connection | None = None
+
+    def connect(self) -> KacheClient:
+        """Establish a dedicated connection (used with context manager)."""
+        self._conn = self._pool.get_connection()
+        return self
+
+    def close(self) -> None:
+        """Release the dedicated connection back to the pool."""
+        if self._conn is not None:
+            self._pool.release_connection(self._conn)
+            self._conn = None
+
+    def disconnect_all(self) -> None:
+        """Close all connections in the pool."""
+        self._pool.disconnect_all()
+
+    def __enter__(self) -> KacheClient:
+        return self.connect()
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+    # ── Internal Helpers ──────────────────────────────────────────────────
+
+    def _get_conn(self) -> Connection:
+        """Get the active connection or checkout from pool."""
+        if self._conn is not None:
+            return self._conn
+        # For non-context-manager usage: get a connection per-call.
+        return self._pool.get_connection()
+
+    def _release_conn(self, conn: Connection) -> None:
+        """Release if it was a per-call checkout."""
+        if self._conn is None:
+            self._pool.release_connection(conn)
+
+    def _execute(self, *args: str | bytes) -> RespValue:
+        """Execute a single command and return the response."""
+        conn = self._get_conn()
+        try:
+            conn.send_command(*args)
+            return conn.read_response()
+        except Exception:
+            # On error, discard the connection.
+            conn.disconnect()
+            raise
+        finally:
+            self._release_conn(conn)
+
+    # ── Redis-Compatible Commands ─────────────────────────────────────────
+
+    def ping(self, message: str | None = None) -> str:
+        """Send ``PING`` and return ``PONG`` or the echoed message.
+
+        Parameters
+        ----------
+        message : str | None
+            Optional message to echo back.
+        """
+        args: list[str | bytes] = ["PING"]
+        if message is not None:
+            args.append(message)
+        result = self._execute(*args)
+        return str(result) if result is not None else "PONG"
+
+    def get(self, key: str | bytes) -> bytes | str | None:
+        """Retrieve the value for *key*.
+
+        Returns ``None`` if the key does not exist or has expired.
+        """
+        return self._execute("GET", key)  # type: ignore[return-value]
+
+    def set(
+        self,
+        key: str | bytes,
+        value: str | bytes,
+        *,
+        ex: int | None = None,
+        px: int | None = None,
+    ) -> bool:
+        """Store *value* under *key* with an optional TTL.
+
+        Parameters
+        ----------
+        key : str | bytes
+            The cache key.
+        value : str | bytes
+            The value to store (binary-safe, up to 2 MB).
+        ex : int | None
+            Expiration time in **seconds**.
+        px : int | None
+            Expiration time in **milliseconds**.
+
+        Returns
+        -------
+        bool
+            ``True`` if the server acknowledged with ``OK``.
+        """
+        args: list[str | bytes] = ["SET", key, value]
+        if ex is not None:
+            args.extend(["EX", str(ex)])
+        elif px is not None:
+            args.extend(["PX", str(px)])
+        result = self._execute(*args)
+        return result == "OK"
+
+    def mget(self, *keys: str | bytes) -> list[bytes | str | None]:
+        """Batch-retrieve values for multiple keys in a single round-trip.
+
+        Returns a list of values (or ``None`` for missing/expired keys)
+        in the same order as the input keys.
+        """
+        if not keys:
+            return []
+        result = self._execute("MGET", *keys)
+        return result if isinstance(result, list) else []  # type: ignore[return-value]
+
+    def delete(self, *keys: str | bytes) -> int:
+        """Delete one or more keys.
+
+        Returns the number of keys that were actually removed.
+        """
+        if not keys:
+            return 0
+        result = self._execute("DEL", *keys)
+        return int(result) if isinstance(result, int) else 0
+
+    def exists(self, *keys: str | bytes) -> int:
+        """Check existence of one or more keys.
+
+        Returns the count of keys that exist and have not expired.
+        """
+        if not keys:
+            return 0
+        result = self._execute("EXISTS", *keys)
+        return int(result) if isinstance(result, int) else 0
+
+    # ── Pipeline ──────────────────────────────────────────────────────────
+
+    def pipeline(self) -> Pipeline:
+        """Create a pipeline for batching multiple commands.
+
+        Usage::
+
+            pipe = client.pipeline()
+            pipe.set("a", "1")
+            pipe.get("a")
+            results = pipe.execute()
+        """
+        conn = self._get_conn()
+        return Pipeline(conn)

kachedb/connection.py

@@ -0,0 +1,143 @@
+"""
+Low-level TCP connection to a KacheDB server.
+
+Manages the socket lifecycle, RESP encoding, and buffered response reading.
+Used internally by :class:`~kachedb.client.KacheClient` and the connection pool.
+"""
+
+from __future__ import annotations
+
+import builtins
+import contextlib
+import socket
+
+from .exceptions import ConnectionError, TimeoutError
+from .resp import RespReader, RespValue, encode_command
+
+
+class Connection:
+    """A single TCP connection to a KacheDB daemon.
+
+    Parameters
+    ----------
+    host : str
+        Server hostname or IP address.
+    port : int
+        Server TCP port.
+    socket_timeout : float | None
+        Socket timeout in seconds.  ``None`` for blocking without timeout.
+    decode_responses : bool
+        If ``True``, decode byte responses to UTF-8 strings.
+    """
+
+    __slots__ = (
+        "_reader",
+        "_sock",
+        "decode_responses",
+        "host",
+        "port",
+        "socket_timeout",
+    )
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 6379,
+        *,
+        socket_timeout: float | None = 5.0,
+        decode_responses: bool = False,
+    ) -> None:
+        self.host = host
+        self.port = port
+        self.socket_timeout = socket_timeout
+        self.decode_responses = decode_responses
+        self._sock: socket.socket | None = None
+        self._reader: RespReader | None = None
+
+    @property
+    def is_connected(self) -> bool:
+        """Return ``True`` if the socket is open."""
+        return self._sock is not None
+
+    def connect(self) -> None:
+        """Establish TCP connection to the KacheDB server."""
+        if self._sock is not None:
+            return
+
+        try:
+            sock = socket.create_connection(
+                (self.host, self.port),
+                timeout=self.socket_timeout,
+            )
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+            self._sock = sock
+            self._reader = RespReader(sock)
+        except OSError as exc:
+            raise ConnectionError(
+                f"Failed to connect to KacheDB at {self.host}:{self.port}: {exc}"
+            ) from exc
+
+    def disconnect(self) -> None:
+        """Close the TCP connection."""
+        if self._sock is not None:
+            with contextlib.suppress(OSError):
+                self._sock.close()
+            self._sock = None
+            self._reader = None
+
+    def send_command(self, *args: str | bytes) -> None:
+        """Encode and send a single RESP command."""
+        if self._sock is None:
+            raise ConnectionError("Not connected to KacheDB")
+
+        data = encode_command(list(args))
+        try:
+            self._sock.sendall(data)
+        except builtins.TimeoutError as exc:
+            self.disconnect()
+            raise TimeoutError(f"Timeout sending command to KacheDB: {exc}") from exc
+        except OSError as exc:
+            self.disconnect()
+            raise ConnectionError(f"Error sending command to KacheDB: {exc}") from exc
+
+    def send_packed(self, data: bytes) -> None:
+        """Send pre-encoded bytes (used by pipeline)."""
+        if self._sock is None:
+            raise ConnectionError("Not connected to KacheDB")
+
+        try:
+            self._sock.sendall(data)
+        except builtins.TimeoutError as exc:
+            self.disconnect()
+            raise TimeoutError(f"Timeout sending data to KacheDB: {exc}") from exc
+        except OSError as exc:
+            self.disconnect()
+            raise ConnectionError(f"Error sending data to KacheDB: {exc}") from exc
+
+    def read_response(self) -> RespValue:
+        """Read and decode one RESP response."""
+        if self._reader is None:
+            raise ConnectionError("Not connected to KacheDB")
+
+        try:
+            response = self._reader.read_response()
+        except builtins.TimeoutError as exc:
+            self.disconnect()
+            raise TimeoutError(f"Timeout reading response from KacheDB: {exc}") from exc
+        except OSError as exc:
+            self.disconnect()
+            raise ConnectionError(f"Error reading response from KacheDB: {exc}") from exc
+
+        if self.decode_responses and isinstance(response, bytes):
+            return response.decode("utf-8")
+
+        return response
+
+    def check_health(self) -> bool:
+        """Send a PING and verify PONG response.  Returns ``False`` on any error."""
+        try:
+            self.send_command("PING")
+            response = self.read_response()
+            return response == "PONG"
+        except Exception:
+            return False