prismdb 0.1.0__py3-none-any.whl

prismdb/__init__.py

@@ -0,0 +1,43 @@
+"""prismdb — a pure-Python client for PrismDB over the binary wire protocol
+(``docs/specs/wire-protocol.md``). No native build, no C extensions.
+
+    from prismdb import Client, Q, U
+
+    with Client.connect(host="127.0.0.1", port=4444, username="admin", password="admin") as db:
+        db.sql("CREATE TABLE users (id BIGINT PRIMARY KEY, name TEXT)")
+        db.sql("INSERT INTO users VALUES (1, 'alice')")
+        print(db.sql("SELECT * FROM users").rows)
+"""
+
+from .client import Client, SqlResult
+from .document import Document
+from .errors import ErrorCode, ErrorInfo, PrismError, PrismServerError, ProtocolError
+from .query import DocQuery, Q
+from .update import DocUpdate, U
+from .value import TAG, ObjectId, Typed, Value, float64, int32, int64, timestamp
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Client",
+    "SqlResult",
+    "Document",
+    "Q",
+    "DocQuery",
+    "U",
+    "DocUpdate",
+    "ObjectId",
+    "Typed",
+    "Value",
+    "TAG",
+    "int32",
+    "int64",
+    "float64",
+    "timestamp",
+    "PrismError",
+    "PrismServerError",
+    "ProtocolError",
+    "ErrorInfo",
+    "ErrorCode",
+    "__version__",
+]

prismdb/_codec.py

@@ -0,0 +1,177 @@
+"""Low-level binary codec: a growable little-endian ``Writer``, a bounds-checked
+``Reader``, and the length-prefixed frame helpers. The byte layouts mirror
+``crates/prism-protocol/src/codec.rs`` exactly (all multi-byte integers LE)."""
+
+from __future__ import annotations
+
+import struct
+
+from .errors import ProtocolError
+
+_U64_MASK = (1 << 64) - 1
+_U128_MASK = (1 << 128) - 1
+
+
+class Writer:
+    """A growable little-endian writer over a bytearray."""
+
+    __slots__ = ("_buf",)
+
+    def __init__(self) -> None:
+        self._buf = bytearray()
+
+    def u8(self, v: int) -> None:
+        self._buf.append(v & 0xFF)
+
+    def u16(self, v: int) -> None:
+        self._buf += struct.pack("<H", v & 0xFFFF)
+
+    def u32(self, v: int) -> None:
+        self._buf += struct.pack("<I", v & 0xFFFFFFFF)
+
+    def i32(self, v: int) -> None:
+        self._buf += struct.pack("<i", _to_signed(v, 32))
+
+    def u64(self, v: int) -> None:
+        self._buf += struct.pack("<Q", v & _U64_MASK)
+
+    def i64(self, v: int) -> None:
+        self._buf += struct.pack("<q", _to_signed(v, 64))
+
+    def f64(self, v: float) -> None:
+        self._buf += struct.pack("<d", v)
+
+    def u128(self, v: int) -> None:
+        """A 128-bit unsigned integer as 16 little-endian bytes."""
+        x = v & _U128_MASK
+        self._buf += struct.pack("<QQ", x & _U64_MASK, (x >> 64) & _U64_MASK)
+
+    def raw(self, b: bytes) -> None:
+        self._buf += b
+
+    def str_u16(self, s: str) -> None:
+        """A UTF-8 string with a u16 length prefix."""
+        b = s.encode("utf-8")
+        self.u16(len(b))
+        self._buf += b
+
+    def str_u32(self, s: str) -> None:
+        """A UTF-8 string with a u32 length prefix."""
+        b = s.encode("utf-8")
+        self.u32(len(b))
+        self._buf += b
+
+    def bytes_u16(self, b: bytes) -> None:
+        """A byte string with a u16 length prefix."""
+        self.u16(len(b))
+        self._buf += b
+
+    def bytes_u32(self, b: bytes) -> None:
+        """A byte string with a u32 length prefix."""
+        self.u32(len(b))
+        self._buf += b
+
+    def out(self) -> bytes:
+        return bytes(self._buf)
+
+
+class Reader:
+    """A bounds-checked little-endian reader over a bytes-like object."""
+
+    __slots__ = ("_buf", "_p")
+
+    def __init__(self, buf: bytes) -> None:
+        self._buf = buf
+        self._p = 0
+
+    def _need(self, n: int) -> None:
+        if self._p + n > len(self._buf):
+            raise ProtocolError(f"truncated: need {n} bytes at offset {self._p}")
+
+    def u8(self) -> int:
+        self._need(1)
+        v = self._buf[self._p]
+        self._p += 1
+        return v
+
+    def u16(self) -> int:
+        self._need(2)
+        v = struct.unpack_from("<H", self._buf, self._p)[0]
+        self._p += 2
+        return v
+
+    def u32(self) -> int:
+        self._need(4)
+        v = struct.unpack_from("<I", self._buf, self._p)[0]
+        self._p += 4
+        return v
+
+    def i32(self) -> int:
+        self._need(4)
+        v = struct.unpack_from("<i", self._buf, self._p)[0]
+        self._p += 4
+        return v
+
+    def u64(self) -> int:
+        self._need(8)
+        v = struct.unpack_from("<Q", self._buf, self._p)[0]
+        self._p += 8
+        return v
+
+    def i64(self) -> int:
+        self._need(8)
+        v = struct.unpack_from("<q", self._buf, self._p)[0]
+        self._p += 8
+        return v
+
+    def f64(self) -> float:
+        self._need(8)
+        v = struct.unpack_from("<d", self._buf, self._p)[0]
+        self._p += 8
+        return v
+
+    def u128(self) -> int:
+        self._need(16)
+        lo, hi = struct.unpack_from("<QQ", self._buf, self._p)
+        self._p += 16
+        return (hi << 64) | lo
+
+    def raw(self, n: int) -> bytes:
+        self._need(n)
+        s = self._buf[self._p : self._p + n]
+        self._p += n
+        return bytes(s)
+
+    def str_u16(self) -> str:
+        return self.raw(self.u16()).decode("utf-8")
+
+    def str_u32(self) -> str:
+        return self.raw(self.u32()).decode("utf-8")
+
+    def bytes_u16(self) -> bytes:
+        return self.raw(self.u16())
+
+    def bytes_u32(self) -> bytes:
+        return self.raw(self.u32())
+
+    def remaining(self) -> int:
+        return len(self._buf) - self._p
+
+    def expect_end(self) -> None:
+        """Raise unless every byte has been consumed."""
+        if self.remaining() != 0:
+            raise ProtocolError(f"{self.remaining()} trailing byte(s) after message")
+
+
+def _to_signed(v: int, bits: int) -> int:
+    """Wrap ``v`` into a signed ``bits``-wide integer (two's complement)."""
+    mask = (1 << bits) - 1
+    v &= mask
+    if v >= 1 << (bits - 1):
+        v -= 1 << bits
+    return v
+
+
+def frame_encode(payload: bytes) -> bytes:
+    """Wrap a payload in a ``[len:u32][payload]`` frame."""
+    return struct.pack("<I", len(payload)) + payload

prismdb/client.py

@@ -0,0 +1,289 @@
+"""The high-level client: connect + handshake, then SQL / KV / document calls
+and transaction control. One client owns one connection = one server session,
+so a ``begin()`` … ``commit()`` brackets the calls in between."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Union
+
+from .connection import Connection, NoticeHandler, TlsArg
+from .document import Document, decode_document, encode_document
+from .errors import ErrorInfo, PrismServerError, ProtocolError
+from .messages import (
+    AUTH_PASSWORD,
+    FEATURE_CONNECT_DB,
+    TXN_READ_ONLY,
+    TXN_READ_WRITE,
+    AuthAck,
+    ColumnDesc,
+    DocResultMsg,
+    HelloAck,
+    KvResultMsg,
+    Pong,
+    SqlResultMsg,
+    TxnAck,
+    abort_body,
+    auth_body,
+    begin_body,
+    commit_body,
+    doc_body,
+    doc_insert_many_body,
+    hello_body,
+    kv_delete_body,
+    kv_get_body,
+    kv_put_body,
+    ping_body,
+    sql_body,
+)
+from .query import DocQuery, Q, encode_doc_query
+from .update import DocUpdate, encode_doc_update
+from .value import ObjectId, Value
+
+_PROTOCOL_VERSION = 1
+_EMPTY = b""
+
+BytesLike = Union[str, bytes, bytearray]
+
+
+@dataclass
+class SqlResult:
+    """A SQL result set. ``rows`` are keyed by column name; ``raw`` keeps cell order."""
+
+    columns: List[ColumnDesc]
+    rows: List[Dict[str, Value]]
+    raw: List[List[Value]]
+    affected_rows: int
+
+
+def _fail(error: Optional[ErrorInfo]) -> "None":
+    raise PrismServerError(
+        error or ErrorInfo(code=0, message="server error", sqlstate="XX000")
+    )
+
+
+def _bytes(v: BytesLike) -> bytes:
+    return v.encode("utf-8") if isinstance(v, str) else bytes(v)
+
+
+class KvSurface:
+    """``client.kv`` — namespaced key/value operations."""
+
+    def __init__(self, client: "Client") -> None:
+        self._c = client
+
+    def get(self, namespace: str, key: BytesLike) -> Optional[bytes]:
+        reply = self._c._kv_reply(*kv_get_body(namespace, _bytes(key)))
+        if reply.op != 1:
+            raise ProtocolError("expected a KV get result")
+        return reply.value
+
+    def put(self, namespace: str, key: BytesLike, value: BytesLike) -> None:
+        self._c._kv_reply(*kv_put_body(namespace, _bytes(key), _bytes(value)))
+
+    def delete(self, namespace: str, key: BytesLike) -> None:
+        self._c._kv_reply(*kv_delete_body(namespace, _bytes(key)))
+
+
+class DocSurface:
+    """``client.doc`` — document collection operations."""
+
+    def __init__(self, client: "Client") -> None:
+        self._c = client
+
+    def insert_one(self, collection: str, document: Document) -> ObjectId:
+        reply = self._c._doc_reply(*doc_body(1, collection, [encode_document(document)]))
+        if not reply.inserted_ids:
+            raise ProtocolError("insert returned no _id")
+        return reply.inserted_ids[0]
+
+    def insert_many(self, collection: str, documents: Sequence[Document]) -> List[ObjectId]:
+        blobs = [encode_document(d) for d in documents]
+        reply = self._c._doc_reply(*doc_insert_many_body(collection, blobs))
+        return reply.inserted_ids
+
+    def find(self, collection: str, query: Optional[DocQuery] = None) -> List[Document]:
+        reply = self._c._doc_reply(*doc_body(3, collection, [encode_doc_query(query or Q.all()), _EMPTY]))
+        return [decode_document(d) for d in reply.docs]
+
+    def find_one(self, collection: str, query: Optional[DocQuery] = None) -> Optional[Document]:
+        reply = self._c._doc_reply(*doc_body(4, collection, [encode_doc_query(query or Q.all()), _EMPTY]))
+        return decode_document(reply.docs[0]) if reply.docs else None
+
+    def count(self, collection: str, query: Optional[DocQuery] = None) -> int:
+        reply = self._c._doc_reply(*doc_body(9, collection, [encode_doc_query(query or Q.all()), _EMPTY]))
+        return reply.affected
+
+    def update_one(self, collection: str, query: DocQuery, update: List[DocUpdate]) -> int:
+        reply = self._c._doc_reply(
+            *doc_body(5, collection, [encode_doc_query(query), encode_doc_update(update), _EMPTY])
+        )
+        return reply.affected
+
+    def update_many(self, collection: str, query: DocQuery, update: List[DocUpdate]) -> int:
+        reply = self._c._doc_reply(
+            *doc_body(6, collection, [encode_doc_query(query), encode_doc_update(update), _EMPTY])
+        )
+        return reply.affected
+
+    def delete_one(self, collection: str, query: DocQuery) -> int:
+        reply = self._c._doc_reply(*doc_body(7, collection, [encode_doc_query(query), _EMPTY]))
+        return reply.affected
+
+    def delete_many(self, collection: str, query: DocQuery) -> int:
+        reply = self._c._doc_reply(*doc_body(8, collection, [encode_doc_query(query), _EMPTY]))
+        return reply.affected
+
+
+class Client:
+    """A connected, authenticated Prism session."""
+
+    def __init__(self, conn: Connection) -> None:
+        self._conn = conn
+        self.kv = KvSurface(self)
+        self.doc = DocSurface(self)
+
+    @classmethod
+    def connect(
+        cls,
+        host: str = "127.0.0.1",
+        port: int = 4444,
+        *,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        database: Optional[str] = None,
+        tls: TlsArg = None,
+        server_hostname: Optional[str] = None,
+        connect_timeout: float = 10.0,
+        client_name: str = "prismdb-python",
+        client_version: str = "0.1.0",
+        on_notice: Optional[NoticeHandler] = None,
+    ) -> "Client":
+        """Connect, perform the handshake, and (if ``username`` is set) authenticate."""
+        conn = Connection.connect(
+            host,
+            port,
+            tls=tls,
+            server_hostname=server_hostname,
+            connect_timeout=connect_timeout,
+            on_notice=on_notice,
+        )
+        client = cls(conn)
+        try:
+            connect_db_honored = client._handshake(
+                username, password, database or "", client_name, client_version
+            )
+            # Fall back to `USE` only when the server did not bind the database
+            # in the handshake (an older server without FEATURE_CONNECT_DB).
+            if database and not connect_db_honored:
+                client.sql(f"USE {database}", return_rows=False)
+        except Exception:
+            conn.close()
+            raise
+        return client
+
+    def _handshake(
+        self,
+        username: Optional[str],
+        password: Optional[str],
+        database: str,
+        client_name: str,
+        client_version: str,
+    ) -> bool:
+        features = FEATURE_CONNECT_DB if database else 0
+        ack = self._conn.request(*hello_body(_PROTOCOL_VERSION, client_name, client_version, features, database))
+        if not isinstance(ack, HelloAck):
+            raise ProtocolError("expected HelloAck")
+        if ack.status != 0:
+            _fail(ack.error)
+        connect_db_honored = (ack.features & FEATURE_CONNECT_DB) != 0 and database != ""
+
+        if username is not None:
+            auth_ack = self._conn.request(*auth_body(AUTH_PASSWORD, username, password or ""))
+            if not isinstance(auth_ack, AuthAck):
+                raise ProtocolError("expected AuthAck")
+            if auth_ack.status != 0:
+                _fail(auth_ack.error)
+        return connect_db_honored
+
+    # ---- SQL --------------------------------------------------------------
+
+    def sql(
+        self,
+        text: str,
+        params: Optional[Sequence[Value]] = None,
+        *,
+        return_rows: bool = True,
+    ) -> SqlResult:
+        """Execute a SQL statement. Returns rows for ``SELECT``, counts otherwise."""
+        reply = self._conn.request(*sql_body(text, list(params or []), 1 if return_rows else 0))
+        if not isinstance(reply, SqlResultMsg):
+            raise ProtocolError("expected SqlResult")
+        if reply.status != 0:
+            _fail(reply.error)
+        if reply.more_frames:
+            raise ProtocolError("streamed SQL results are not yet supported")
+        names = [c.name for c in reply.columns]
+        rows = [{names[i]: cell for i, cell in enumerate(cells)} for cells in reply.rows]
+        return SqlResult(reply.columns, rows, reply.rows, reply.affected_rows)
+
+    # ---- transactions -----------------------------------------------------
+
+    def begin(self, mode: str = "read_write") -> int:
+        """Begin a transaction; returns the assigned transaction id."""
+        ack = self._txn(*begin_body(TXN_READ_ONLY if mode == "read_only" else TXN_READ_WRITE))
+        return ack.txn_id
+
+    def commit(self, idempotency_key: int = 0) -> None:
+        """Commit the current transaction (optionally idempotent)."""
+        self._txn(*commit_body(idempotency_key))
+
+    def abort(self) -> None:
+        """Abort the current transaction."""
+        self._txn(*abort_body())
+
+    def _txn(self, type_code: int, body: bytes) -> TxnAck:
+        reply = self._conn.request(type_code, body)
+        if not isinstance(reply, TxnAck):
+            raise ProtocolError("expected TxnAck")
+        if reply.status != 0:
+            _fail(reply.error)
+        return reply
+
+    # ---- misc -------------------------------------------------------------
+
+    def ping(self) -> None:
+        """Round-trip a keep-alive ping."""
+        reply = self._conn.request(*ping_body())
+        if not isinstance(reply, Pong):
+            raise ProtocolError("expected Pong")
+
+    def close(self) -> None:
+        """Close the underlying connection."""
+        self._conn.close()
+
+    def __enter__(self) -> "Client":
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    # ---- internal reply helpers -------------------------------------------
+
+    def _kv_reply(self, type_code: int, body: bytes) -> KvResultMsg:
+        reply = self._conn.request(type_code, body)
+        if not isinstance(reply, KvResultMsg):
+            raise ProtocolError("expected KvResult")
+        if reply.status != 0:
+            _fail(reply.error)
+        return reply
+
+    def _doc_reply(self, type_code: int, body: bytes) -> DocResultMsg:
+        reply = self._conn.request(type_code, body)
+        if not isinstance(reply, DocResultMsg):
+            raise ProtocolError("expected DocResult")
+        if reply.status != 0:
+            _fail(reply.error)
+        if reply.more_frames:
+            raise ProtocolError("streamed document results are not yet supported")
+        return reply

prismdb/connection.py

@@ -0,0 +1,116 @@
+"""The transport: a TCP (optionally TLS) socket that frames outgoing messages,
+reads back full frames, and matches each reply to its request by the echoed
+``request_id``. Server-initiated notices (request_id 0) go to a handler.
+
+The connection is synchronous: each :meth:`request` writes a frame and blocks
+until the matching reply arrives, dispatching any notices seen in between."""
+
+from __future__ import annotations
+
+import socket
+import ssl
+from typing import Callable, Optional, Tuple, Union
+
+from ._codec import frame_encode
+from .errors import ProtocolError
+from .messages import Notice, decode_packet, encode_packet
+
+NoticeHandler = Callable[[Notice], None]
+# ``tls``: False/None = plaintext; True = default client context; an
+# ``ssl.SSLContext`` = use it as-is.
+TlsArg = Union[bool, ssl.SSLContext, None]
+
+
+class Connection:
+    """A framed, request/reply socket connection to a Prism server."""
+
+    def __init__(self, sock: socket.socket, on_notice: Optional[NoticeHandler] = None) -> None:
+        self._sock = sock
+        self._on_notice = on_notice
+        self._next_id = 1
+        self._inbound = bytearray()
+        self._closed: Optional[Exception] = None
+
+    @classmethod
+    def connect(
+        cls,
+        host: str = "127.0.0.1",
+        port: int = 4444,
+        *,
+        tls: TlsArg = None,
+        server_hostname: Optional[str] = None,
+        connect_timeout: float = 10.0,
+        on_notice: Optional[NoticeHandler] = None,
+    ) -> "Connection":
+        """Open a connection (TCP, or TLS when ``tls`` is set)."""
+        raw = socket.create_connection((host, port), timeout=connect_timeout)
+        raw.settimeout(None)
+        raw.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        sock: socket.socket = raw
+        if tls:
+            ctx = tls if isinstance(tls, ssl.SSLContext) else ssl.create_default_context()
+            sock = ctx.wrap_socket(raw, server_hostname=server_hostname or host)
+        return cls(sock, on_notice)
+
+    def request(self, type_code: int, body: bytes) -> object:
+        """Send a client message and return the matching reply message."""
+        if self._closed is not None:
+            raise self._closed
+        request_id = self._next_id
+        self._next_id = 1 if self._next_id >= 0xFFFFFFFF else self._next_id + 1
+        try:
+            self._sock.sendall(frame_encode(encode_packet(request_id, type_code, body)))
+        except OSError as e:
+            self._fail(ProtocolError(f"send failed: {e}"))
+            raise self._closed  # type: ignore[misc]
+
+        while True:
+            payload = self._read_frame()
+            packet = decode_packet(payload)
+            if isinstance(packet.message, Notice):
+                if self._on_notice is not None:
+                    self._on_notice(packet.message)
+                continue
+            if packet.request_id == request_id:
+                return packet.message
+            # An unmatched reply (e.g. a late response) is ignored.
+
+    def close(self) -> None:
+        """Close the connection. Further use raises."""
+        self._fail(ProtocolError("connection closed by client"))
+        try:
+            self._sock.close()
+        except OSError:
+            pass
+
+    # -- internals ----------------------------------------------------------
+
+    def _read_frame(self) -> bytes:
+        header = self._read_exact(4)
+        length = int.from_bytes(header, "little")
+        return self._read_exact(length)
+
+    def _read_exact(self, n: int) -> bytes:
+        # Drain anything already buffered from a previous read first.
+        while len(self._inbound) < n:
+            try:
+                chunk = self._sock.recv(65536)
+            except OSError as e:
+                self._fail(ProtocolError(f"connection closed by server: {e}"))
+                raise self._closed  # type: ignore[misc]
+            if not chunk:
+                self._fail(ProtocolError("connection closed by server"))
+                raise self._closed  # type: ignore[misc]
+            self._inbound += chunk
+        out = bytes(self._inbound[:n])
+        del self._inbound[:n]
+        return out
+
+    def _fail(self, err: Exception) -> None:
+        if self._closed is None:
+            self._closed = err
+
+
+def split_tls_args(tls: TlsArg) -> Tuple[TlsArg, Optional[str]]:  # pragma: no cover
+    """Reserved hook for richer TLS option parsing; identity for now."""
+    return tls, None

prismdb/document.py

@@ -0,0 +1,50 @@
+"""The document tagged-binary codec.
+
+Mirrors ``crates/prism-doc/src/value.rs`` (``Document::encode``/``decode``). A
+document is ``[total:u32][count:u16]`` followed by, per field,
+``[tag:u8][nameLen:u16][name][value bytes]``. Field value bytes use the same
+encoding as scalar values, except documents have no Binary type."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+from ._codec import Reader, Writer
+from .errors import ProtocolError
+from .value import TAG, Value, decode_untagged, encode_untagged, tag_of
+
+# A document is a plain dict; field insertion order is preserved.
+Document = Dict[str, Value]
+
+
+def encode_document(doc: Document) -> bytes:
+    """Encode a document to its tagged-binary payload."""
+    body = Writer()
+    if len(doc) > 0xFFFF:
+        raise ProtocolError("too many document fields")
+    body.u16(len(doc))
+    for name, value in doc.items():
+        tag = tag_of(value)
+        if tag == TAG.BINARY:
+            raise ProtocolError(f'field "{name}": binary values are not supported in documents')
+        body.u8(tag)
+        body.str_u16(name)
+        encode_untagged(body, tag, value)
+    inner = body.out()
+    out = Writer()
+    out.u32(4 + len(inner))  # total length, including this u32
+    out.raw(inner)
+    return out.out()
+
+
+def decode_document(raw: bytes) -> Document:
+    """Decode a document from its tagged-binary payload."""
+    r = Reader(raw)
+    r.u32()  # total length (redundant with the frame's blob length)
+    count = r.u16()
+    doc: Document = {}
+    for _ in range(count):
+        tag = r.u8()
+        name = r.str_u16()
+        doc[name] = decode_untagged(r, tag)
+    return doc