bladekit 4.5.0__py3-none-any.whl

bladekit/_queue.py

@@ -0,0 +1,137 @@
+"""bladekit.Queue — distributed FIFO queue.
+
+Persistent, backed by Postgres on the control plane. Used for producer/consumer
+coordination between functions or between local scripts and remote workers.
+
+Protocol (conventions.md §14.5, map-protocol.md is a SEPARATE thing):
+- put(value): append an item to the tail of the queue.
+- get(timeout=0): claim-and-delete the head item atomically
+  (FOR UPDATE SKIP LOCKED). Two concurrent get() calls never return the same
+  item. Returns None if the queue is empty and timeout=0. Blocks up to timeout
+  seconds (long-poll on the server) if timeout > 0.
+- len(): current item count (non-transactional, best-effort).
+- from_name(name): get-or-create by name (global namespace, RBAC planned).
+
+Limits (conventions.md §6):
+- Up to 100,000 items per queue.
+- Blocking get: timeout 0–600 s.
+- Item payload: ≤ 1 MiB inline; object store beyond.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class Queue:
+    """A distributed FIFO queue persistent in the AI Fabric control plane.
+
+    Usage::
+
+        q = bladekit.Queue.from_name("jobs")
+        q.put({"url": "https://example.com"})
+        item = q.get(timeout=30)   # blocks up to 30s
+
+    Items are serialized with cloudpickle (same format as function arguments).
+    The queue is identified by name; ``from_name`` is idempotent (get-or-create).
+    """
+
+    def __init__(self, name: str, client: Any) -> None:
+        self._name = name
+        self._client = client
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"bladekit.Queue(name={self._name!r})"
+
+    @classmethod
+    def from_name(cls, name: str, *, create_if_missing: bool = False) -> "Queue":
+        """Reference a queue by name (get-or-create on first use).
+
+        No round-trip here: the queue is created lazily by the server on the
+        first put/get (openapi.yaml: "get-or-create by name on first use"). Two
+        references with the same name address the same queue. ``create_if_missing``
+        is accepted for parity with ``Volume.from_name`` (and Modal) but is a
+        no-op — queues are always get-or-create on first use.
+        """
+        _ = create_if_missing
+        from ._client import Client
+        from .config import load_config
+        from . import exceptions as exc_mod
+
+        cfg = load_config()
+        if not cfg.token:
+            raise exc_mod.AuthenticationError(
+                "UNAUTHENTICATED",
+                "no token configured: run `blade token set <token>` or set FABRIC_TOKEN",
+            )
+        return cls(name, Client(cfg.api_url, cfg.token))
+
+    @classmethod
+    def _from_client(cls, name: str, client: Any) -> "Queue":
+        """Create a Queue object backed by an existing client (for tests)."""
+        return cls(name, client)
+
+    # -- operations -------------------------------------------------------
+
+    def put(self, value: Any) -> None:
+        """Append an item to the tail of the queue.
+
+        The item is serialized with cloudpickle. Raises PayloadTooLargeError
+        if the serialized item exceeds 1 MiB (QUEUE_ITEM_TOO_LARGE).
+        """
+        from ._serialize import dumps_args, make_payload_in
+
+        raw = dumps_args((value,), {})
+        payload_in, object_bytes = make_payload_in(raw)
+        if object_bytes is not None:
+            decl = self._client.create_payload(
+                payload_in["size_bytes"], payload_in["sha256"]
+            )
+            self._client.upload_to_presigned(
+                decl["put_url"], object_bytes, "application/octet-stream"
+            )
+            payload_in["payload_id"] = decl["payload_id"]
+        self._client.queue_put(self._name, payload_in)
+
+    def get(self, *, timeout: int = 0) -> Optional[Any]:
+        """Claim and remove the head item.
+
+        Returns the item value, or None if the queue is empty and timeout=0.
+        With timeout > 0, the call blocks on the server for up to ``timeout``
+        seconds (long-poll, max 600 s). Returns None on timeout.
+
+        Two concurrent get() calls never return the same item.
+        """
+        from ._serialize import decode_payload_out, loads
+
+        resp = self._client.queue_get(self._name, timeout=timeout)
+        # QueueGetResponse: {item: PayloadOut|null, item_id}. item is the payload
+        # directly (not wrapped); null = empty.
+        payload = resp.get("item")
+        if payload is None:
+            return None
+        raw = decode_payload_out(payload, fetch=self._client.download_presigned)
+        # put() serialized (value,) as args; unwrap the tuple.
+        result = loads(raw)
+        if isinstance(result, tuple) and len(result) == 2:
+            args, kwargs = result
+            if isinstance(args, tuple) and len(args) == 1 and not kwargs:
+                return args[0]
+        return result
+
+    def len(self) -> int:
+        """Return the approximate number of items currently in the queue.
+
+        The count is non-transactional (another consumer may have claimed items
+        between the request and this return value).
+        """
+        resp = self._client.get_queue(self._name)
+        return int(resp.get("len", 0))
+
+    def delete(self) -> None:
+        """Delete the queue and all its items."""
+        self._client.delete_queue(self._name)