flo-python 0.1.0.dev4__py3-none-any.whl → 0.1.1__py3-none-any.whl

flo/kv_txn.py

@@ -0,0 +1,340 @@
+"""Flo KV Per-Shard Transactions
+
+Transactions are pinned to a single partition (chosen by the routing key
+passed to ``begin``). Every key written or read inside the transaction must
+hash to the same partition; otherwise the server returns a
+"kv_txn_cross_shard" error.
+
+Caps (server-enforced):
+    * 256 ops per transaction
+    * 1 MiB total payload across buffered writes
+
+The following operations are NOT supported inside a transaction and raise
+:class:`TxnUnsupportedOpError` without a server round-trip: ``scan``,
+``mget``, ``json_get``, ``json_set``, ``json_del``, ``history``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import struct
+from typing import TYPE_CHECKING, Any
+
+from .types import (
+    DeleteOptions,
+    GetResult,
+    KVCommitResult,
+    OpCode,
+    OptionTag,
+    PutOptions,
+    PutResult,
+)
+from .wire import OptionsBuilder
+
+if TYPE_CHECKING:
+    from .client import FloClient
+
+
+class TxnUnsupportedOpError(Exception):
+    """Raised when an operation is not allowed inside a KV transaction."""
+
+
+class TxnFinishedError(Exception):
+    """Raised when an operation is attempted on a closed transaction."""
+
+
+class Transaction:
+    """Per-shard KV transaction handle.
+
+    Operations are buffered on the server's pinned shard until :meth:`commit`
+    or :meth:`rollback` is called. Use as an async context manager for
+    automatic rollback on exception::
+
+        async with await client.kv.begin("user:123") as txn:
+            await txn.put("user:123:name", b"Jane")
+            await txn.incr("user:123:visits", 1)
+            result = await txn.commit()
+            print(result.commit_index)
+    """
+
+    def __init__(
+        self,
+        client: FloClient,
+        namespace: str,
+        routing_key: str,
+        txn_id: int,
+        pinned_hash: int,
+    ):
+        self._client = client
+        self._namespace = namespace
+        self._routing_key = routing_key
+        self._txn_id = txn_id
+        self._pinned_hash = pinned_hash
+        self._done = False
+
+    @property
+    def id(self) -> int:
+        """The server-assigned transaction id."""
+        return self._txn_id
+
+    @property
+    def pinned_hash(self) -> int:
+        """The partition hash this transaction is bound to."""
+        return self._pinned_hash
+
+    def _txn_options(self) -> OptionsBuilder:
+        b = OptionsBuilder()
+        if self._routing_key:
+            b.add_bytes(OptionTag.ROUTING_KEY, self._routing_key.encode("utf-8"))
+        b.add_u64(OptionTag.TXN_ID, self._txn_id)
+        return b
+
+    def _check_alive(self) -> None:
+        if self._done:
+            raise TxnFinishedError("transaction already committed or rolled back")
+
+    async def put(
+        self,
+        key: str | bytes,
+        value: bytes,
+        options: PutOptions | None = None,
+    ) -> PutResult:
+        """Buffer a put inside the transaction."""
+        self._check_alive()
+        opts = options or PutOptions()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        if opts.ttl_seconds is not None:
+            builder.add_u64(OptionTag.TTL_SECONDS, opts.ttl_seconds)
+        if opts.cas_version is not None:
+            builder.add_u64(OptionTag.CAS_VERSION, opts.cas_version)
+        if opts.if_not_exists:
+            builder.add_flag(OptionTag.IF_NOT_EXISTS)
+        if opts.if_exists:
+            builder.add_flag(OptionTag.IF_EXISTS)
+        response = await self._client._send_and_check(
+            OpCode.KV_PUT,
+            self._namespace,
+            key_bytes,
+            value,
+            builder.build(),
+        )
+        if not response.data or len(response.data) < 8:
+            return PutResult(version=0)
+        return PutResult(version=struct.unpack("<Q", response.data[:8])[0])
+
+    async def get(self, key: str | bytes) -> GetResult | None:
+        """Read a key inside the transaction (sees buffered writes)."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        response = await self._client._send_and_check(
+            OpCode.KV_GET,
+            self._namespace,
+            key_bytes,
+            b"",
+            builder.build(),
+            allow_not_found=True,
+        )
+        if response.is_not_found():
+            return None
+        if not response.data or len(response.data) < 8:
+            return GetResult(value=b"", version=0)
+        version = struct.unpack("<Q", response.data[:8])[0]
+        return GetResult(value=bytes(response.data[8:]), version=version)
+
+    async def delete(
+        self,
+        key: str | bytes,
+        options: DeleteOptions | None = None,
+    ) -> None:
+        """Buffer a delete inside the transaction."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        await self._client._send_and_check(
+            OpCode.KV_DELETE,
+            self._namespace,
+            key_bytes,
+            b"",
+            builder.build(),
+            allow_not_found=True,
+        )
+
+    async def incr(self, key: str | bytes, delta: int = 1) -> int:
+        """Buffer an atomic counter increment inside the transaction."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        value = struct.pack("<q", delta)
+        response = await self._client._send_and_check(
+            OpCode.KV_INCR,
+            self._namespace,
+            key_bytes,
+            value,
+            builder.build(),
+        )
+        if not response.data or len(response.data) < 8:
+            return 0
+        return int(struct.unpack("<q", response.data[:8])[0])
+
+    async def touch(self, key: str | bytes, ttl_seconds: int) -> None:
+        """Update the TTL on an existing key inside the transaction."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        await self._client._send_and_check(
+            OpCode.KV_TOUCH,
+            self._namespace,
+            key_bytes,
+            struct.pack("<Q", ttl_seconds),
+            builder.build(),
+        )
+
+    async def persist(self, key: str | bytes) -> None:
+        """Remove the TTL on a key inside the transaction."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        await self._client._send_and_check(
+            OpCode.KV_PERSIST,
+            self._namespace,
+            key_bytes,
+            b"",
+            builder.build(),
+        )
+
+    async def exists(self, key: str | bytes) -> bool:
+        """Check key existence inside the transaction."""
+        self._check_alive()
+        key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+        builder = self._txn_options()
+        response = await self._client._send_and_check(
+            OpCode.KV_EXISTS,
+            self._namespace,
+            key_bytes,
+            b"",
+            builder.build(),
+        )
+        # Wire body: [version:u64 LE][1 byte 0/1]
+        if not response.data or len(response.data) < 9:
+            return False
+        return response.data[8] == 1
+
+    # ── Disallowed inside a transaction ───────────────────────────────
+
+    async def scan(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("scan is not supported inside a KV transaction")
+
+    async def mget(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("mget is not supported inside a KV transaction")
+
+    async def json_get(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("json_get is not supported inside a KV transaction")
+
+    async def json_set(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("json_set is not supported inside a KV transaction")
+
+    async def json_del(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("json_del is not supported inside a KV transaction")
+
+    async def history(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+        raise TxnUnsupportedOpError("history is not supported inside a KV transaction")
+
+    # ── Lifecycle ──────────────────────────────────────────────────────
+
+    async def commit(self) -> KVCommitResult:
+        """Atomically apply all buffered operations.
+
+        After commit returns the transaction is closed and further operations
+        raise :class:`TxnFinishedError`.
+        """
+        self._check_alive()
+        self._done = True
+        builder = OptionsBuilder()
+        if self._routing_key:
+            builder.add_bytes(OptionTag.ROUTING_KEY, self._routing_key.encode("utf-8"))
+        builder.add_u64(OptionTag.TXN_ID, self._txn_id)
+        response = await self._client._send_and_check(
+            OpCode.KV_COMMIT_TXN,
+            self._namespace,
+            self._routing_key.encode("utf-8"),
+            b"",
+            builder.build(),
+        )
+        # Wire body: [variant:u8=1][commit_index:u64 LE][op_count:u16 LE]
+        if not response.data or len(response.data) < 11:
+            raise RuntimeError(
+                f"flo: short KV commit reply ({len(response.data) if response.data else 0} bytes)"
+            )
+        return KVCommitResult(
+            commit_index=struct.unpack("<Q", response.data[1:9])[0],
+            op_count=struct.unpack("<H", response.data[9:11])[0],
+        )
+
+    async def rollback(self) -> None:
+        """Discard the buffered operations without committing.
+
+        Idempotent: calling rollback after commit (or vice versa) is a no-op.
+        """
+        if self._done:
+            return
+        self._done = True
+        builder = OptionsBuilder()
+        if self._routing_key:
+            builder.add_bytes(OptionTag.ROUTING_KEY, self._routing_key.encode("utf-8"))
+        builder.add_u64(OptionTag.TXN_ID, self._txn_id)
+        await self._client._send_and_check(
+            OpCode.KV_ROLLBACK_TXN,
+            self._namespace,
+            self._routing_key.encode("utf-8"),
+            b"",
+            builder.build(),
+        )
+
+    async def __aenter__(self) -> Transaction:
+        return self
+
+    async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
+        if not self._done:
+            # Best-effort cleanup; surface the original exception.
+            with contextlib.suppress(Exception):
+                await self.rollback()
+
+
+async def begin(
+    client: FloClient,
+    namespace: str,
+    routing_key: str,
+) -> Transaction:
+    """Open a new per-shard KV transaction."""
+    builder = OptionsBuilder()
+    if routing_key:
+        builder.add_bytes(OptionTag.ROUTING_KEY, routing_key.encode("utf-8"))
+    response = await client._send_and_check(
+        OpCode.KV_BEGIN_TXN,
+        namespace,
+        routing_key.encode("utf-8"),
+        b"",
+        builder.build(),
+    )
+    # Wire body: [variant:u8=0][txn_id:u64 LE][pinned_hash:u64 LE]
+    if not response.data or len(response.data) < 17:
+        raise RuntimeError(
+            f"flo: short KV begin reply ({len(response.data) if response.data else 0} bytes)"
+        )
+    return Transaction(
+        client=client,
+        namespace=namespace,
+        routing_key=routing_key,
+        txn_id=struct.unpack("<Q", response.data[1:9])[0],
+        pinned_hash=struct.unpack("<Q", response.data[9:17])[0],
+    )
+
+
+__all__ = [
+    "Transaction",
+    "TxnUnsupportedOpError",
+    "TxnFinishedError",
+    "begin",
+]

flo/streams.py

@@ -9,8 +9,10 @@ from typing import TYPE_CHECKING
 from .types import (
     OpCode,
     OptionTag,
+    PendingEntry,
     StreamAppendOptions,
     StreamAppendResult,
+    StreamClaimResult,
     StreamGroupAckOptions,
     StreamGroupJoinOptions,
     StreamGroupNackOptions,
@@ -24,13 +26,21 @@ from .types import (
 )
 from .wire import (
     OptionsBuilder,
+    build_stream_batch_value,
+    parse_pending_entries,
     parse_stream_append_response,
     parse_stream_info_response,
     parse_stream_read_response,
     serialize_group_ack_value,
+    serialize_group_claim_value,
+    serialize_group_pending_value,
     serialize_group_value,
 )
 
+# u64 max sentinel: group_claim returns StreamID.MAX when the PEL is fully
+# scanned (mirrors Go's math.MaxUint64 check).
+_U64_MAX = 0xFFFFFFFFFFFFFFFF
+
 if TYPE_CHECKING:
     from .client import FloClient
 
@@ -63,12 +73,13 @@ class StreamOperations:
         """
         opts = options or StreamAppendOptions()
         namespace = self._client.get_namespace(opts.namespace)
+        value = build_stream_batch_value(payload, opts.headers)
 
         response = await self._client._send_and_check(
             OpCode.STREAM_APPEND,
             namespace,
             stream.encode("utf-8"),
-            payload,
+            value,
             allow_not_found=True,
         )
 
@@ -294,11 +305,18 @@ class StreamOperations:
         consumer: str,
         options: StreamGroupReadOptions | None = None,
     ) -> StreamReadResult:
-        """Read records from a consumer group.
+        """Read new records from a consumer group, advancing last_delivered_id.
+
+        Records are distributed among consumers in the group; each record is
+        delivered to only one consumer and added to that consumer's Pending
+        Entry List (PEL) until acked. Unacknowledged records are redelivered.
 
-        Records are distributed among consumers in the group. Each record
-        is delivered to only one consumer. Unacknowledged records will be
-        redelivered.
+        Crash recovery: ``group_read`` alone is NOT sufficient. It only returns
+        records past ``last_delivered_id``, so a record delivered-but-unacked at
+        crash time is never re-surfaced by a later ``group_read``. To re-process
+        in-flight work after a reconnect, drain the PEL with :meth:`group_claim`
+        (StreamWorker does this automatically — see
+        ``redeliver_pending_on_reconnect``).
 
         Args:
             stream: Stream name.
@@ -340,6 +358,99 @@ class StreamOperations:
 
         return parse_stream_read_response(response.data)
 
+    async def group_pending(
+        self,
+        stream: str,
+        group: str,
+        consumer: str = "",
+        options: StreamGroupReadOptions | None = None,
+    ) -> list[PendingEntry]:
+        """List a consumer group's pending (delivered-but-unacked) entries.
+
+        If ``consumer`` is non-empty, only that consumer's entries are returned;
+        otherwise the whole group's PEL is returned. (FLO-102)
+
+        Returns:
+            A list of :class:`PendingEntry`.
+        """
+        opts = options or StreamGroupReadOptions()
+        namespace = self._client.get_namespace(opts.namespace)
+
+        value = serialize_group_pending_value(group, consumer)
+
+        response = await self._client._send_and_check(
+            OpCode.STREAM_GROUP_PENDING,
+            namespace,
+            stream.encode("utf-8"),
+            value,
+            allow_not_found=True,
+        )
+
+        return parse_pending_entries(response.data)
+
+    async def group_claim(
+        self,
+        stream: str,
+        group: str,
+        consumer: str,
+        min_idle_ms: int,
+        start_id: StreamID,
+        count: int,
+        options: StreamGroupReadOptions | None = None,
+    ) -> StreamClaimResult:
+        """Claim a page of a consumer group's pending entries for ``consumer``.
+
+        Scans the PEL in StreamID order from ``start_id`` and takes up to
+        ``count`` entries idle for at least ``min_idle_ms``. Returns the claimed
+        records (payload + headers) plus a cursor for the next page. (FLO-102)
+
+        * Drain own pending (reconnect): ``min_idle_ms=0``, ``start_id=StreamID()``.
+        * Steal from idle consumers (rebalance): ``min_idle_ms > 0``.
+
+        Loop until ``result.done`` to fully drain::
+
+            cursor = StreamID()
+            while True:
+                r = await client.stream.group_claim(
+                    stream, group, consumer, 0, cursor, 100)
+                for rec in r.records:
+                    process(rec)
+                if r.done or not r.records:
+                    break
+                cursor = r.next_cursor
+        """
+        opts = options or StreamGroupReadOptions()
+        namespace = self._client.get_namespace(opts.namespace)
+
+        value = serialize_group_claim_value(group, consumer, min_idle_ms, start_id, count)
+
+        response = await self._client._send_and_check(
+            OpCode.STREAM_GROUP_CLAIM,
+            namespace,
+            stream.encode("utf-8"),
+            value,
+            allow_not_found=True,
+        )
+
+        # Response = <records blob> + [next_ts:u64][next_seq:u64] trailer.
+        data = response.data or b""
+        if len(data) < 16:
+            return StreamClaimResult(records=[], next_cursor=StreamID(), done=True)
+
+        cursor_off = len(data) - 16
+        next_ts = int.from_bytes(data[cursor_off : cursor_off + 8], "little")
+        next_seq = int.from_bytes(data[cursor_off + 8 : cursor_off + 16], "little")
+
+        read = parse_stream_read_response(data[:cursor_off])
+
+        # StreamID.MAX (max, max) is the "fully scanned" sentinel.
+        done = next_ts == _U64_MAX and next_seq == _U64_MAX
+        return StreamClaimResult(
+            records=read.records,
+            next_cursor=StreamID(timestamp_ms=next_ts, sequence=next_seq),
+            done=done,
+        )
+
     async def group_ack(
         self,
         stream: str,