leanqueue 0.1.0__cp38-abi3-win_amd64.whl

leanqueuedb.py

@@ -0,0 +1,411 @@
+"""LeanQueue — embedded, durable, crash-resistant message queue.
+
+Thin Python orchestration over the Rust core (``leanqueue.LeanQueueCore``):
+a friendly ``Queue`` object, a ``Message`` wrapper with ack/nack, JSON/text
+helpers, blocking consume loops, and a ``Broker`` that manages many named
+queues under one directory. All durability, crash recovery, and multi-process /
+multi-thread coordination live in Rust; this layer only adds ergonomics.
+
+Quick start
+-----------
+    from leanqueuedb import Queue
+
+    q = Queue("./jobs")
+    q.put(b"hello")
+    q.put_json({"task": "resize", "id": 42})
+
+    for msg in q.consume(timeout=1.0):     # auto-acks each message on success
+        print(msg.text(), msg.json() if msg.is_json() else None)
+
+    # or manual control / at-least-once with redelivery on crash:
+    msg = q.get(block=False)
+    if msg:
+        try:
+            handle(msg)
+            msg.ack()
+        except Exception:
+            msg.nack(requeue=True)
+"""
+
+import atexit
+import os
+import threading
+import time
+import weakref
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+import leanqueue
+
+try:
+    import orjson
+
+    def _json_dumps(obj: Any) -> bytes:
+        return orjson.dumps(obj)
+
+    def _json_loads(data: bytes) -> Any:
+        return orjson.loads(data)
+except ImportError:  # pragma: no cover - orjson is a declared dependency
+    import json
+
+    def _json_dumps(obj: Any) -> bytes:
+        return json.dumps(obj, separators=(",", ":")).encode("utf-8")
+
+    def _json_loads(data: bytes) -> Any:
+        return json.loads(data.decode("utf-8"))
+
+
+__version__ = getattr(leanqueue, "__version__", "0.1.0")
+
+BodyLike = Union[bytes, bytearray, memoryview, str]
+
+
+def _to_bytes(body: BodyLike) -> bytes:
+    if isinstance(body, str):
+        return body.encode("utf-8")
+    if isinstance(body, (bytearray, memoryview)):
+        return bytes(body)
+    if isinstance(body, bytes):
+        return body
+    raise TypeError(
+        f"message body must be bytes/str, got {type(body).__name__}; "
+        "use put_json() for objects"
+    )
+
+
+# Track every live Queue so we can fsync on interpreter shutdown.
+_OPEN_QUEUES: "weakref.WeakSet[Queue]" = weakref.WeakSet()
+
+
+@atexit.register
+def _flush_all_on_exit() -> None:
+    for q in list(_OPEN_QUEUES):
+        try:
+            q.flush()
+        except Exception:
+            pass
+
+
+class Message:
+    """A leased message. Settle it with :meth:`ack` (done) or :meth:`nack`
+    (failed → redeliver or dead-letter). If never settled, the lease expires and
+    the message is redelivered automatically."""
+
+    __slots__ = ("seq", "body", "deliveries", "enqueued_at_ms", "_queue", "_settled")
+
+    def __init__(self, seq: int, body: bytes, deliveries: int, enqueued_at_ms: int, queue: "Queue"):
+        self.seq = seq
+        self.body = body
+        self.deliveries = deliveries
+        self.enqueued_at_ms = enqueued_at_ms
+        self._queue = queue
+        self._settled = False
+
+    @property
+    def enqueued_at(self) -> datetime:
+        return datetime.fromtimestamp(self.enqueued_at_ms / 1000, tz=timezone.utc)
+
+    @property
+    def settled(self) -> bool:
+        return self._settled
+
+    def text(self, encoding: str = "utf-8") -> str:
+        return self.body.decode(encoding)
+
+    def json(self) -> Any:
+        return _json_loads(self.body)
+
+    def is_json(self) -> bool:
+        try:
+            _json_loads(self.body)
+            return True
+        except Exception:
+            return False
+
+    def ack(self) -> bool:
+        """Mark processing complete; the message is removed permanently."""
+        self._settled = True
+        return self._queue.ack(self.seq)
+
+    def nack(self, requeue: bool = True) -> bool:
+        """Mark processing failed. ``requeue=True`` redelivers it later;
+        ``requeue=False`` drops (dead-letters) it."""
+        self._settled = True
+        return self._queue.nack(self.seq, requeue=requeue)
+
+    def extend_lease(self, lease_ms: Optional[int] = None) -> bool:
+        """Renew the lease while still working on a long task."""
+        return self._queue.extend_lease(self.seq, lease_ms=lease_ms)
+
+    def __repr__(self) -> str:
+        preview = self.body[:32]
+        return (
+            f"Message(seq={self.seq}, deliveries={self.deliveries}, "
+            f"body={preview!r}{'...' if len(self.body) > 32 else ''})"
+        )
+
+
+class Queue:
+    """A single durable queue backed by one directory (or ``':memory:'``).
+
+    Parameters
+    ----------
+    path: directory for this queue, or ``":memory:"`` for a non-persistent one.
+    sync_interval_ms: ``0`` (default) fsyncs after every write — strongest
+        durability. A positive value coalesces fsyncs, trading a small crash
+        window for higher ingest throughput.
+    default_lease_ms: how long a dequeued message stays invisible before
+        redelivery if not acked (default 30s).
+    auto_compact_bytes: rewrite the log to reclaim acked space once it grows
+        past this size with mostly-dead content (default 64 MiB; ``0`` disables).
+    """
+
+    def __init__(
+        self,
+        path: str,
+        sync_interval_ms: int = 0,
+        default_lease_ms: int = 30_000,
+        auto_compact_bytes: int = 64 * 1024 * 1024,
+    ):
+        self.path = path
+        self._core = leanqueue.LeanQueueCore(
+            path,
+            sync_interval_ms=sync_interval_ms,
+            default_lease_ms=default_lease_ms,
+            auto_compact_bytes=auto_compact_bytes,
+        )
+        self._closed = False
+        _OPEN_QUEUES.add(self)
+
+    # --- producing ---
+
+    def put(self, body: BodyLike) -> int:
+        """Enqueue one message; returns its sequence id."""
+        return self._core.enqueue(_to_bytes(body))
+
+    def put_json(self, obj: Any) -> int:
+        """Enqueue a JSON-serializable object."""
+        return self._core.enqueue(_json_dumps(obj))
+
+    def put_many(self, bodies: List[BodyLike]) -> List[int]:
+        """Enqueue many messages under a single lock and fsync."""
+        return self._core.enqueue_batch([_to_bytes(b) for b in bodies])
+
+    def put_many_json(self, objs: List[Any]) -> List[int]:
+        return self._core.enqueue_batch([_json_dumps(o) for o in objs])
+
+    # --- consuming ---
+
+    def _wrap(self, raw) -> Optional[Message]:
+        if raw is None:
+            return None
+        seq, body, deliveries, ts = raw
+        return Message(seq, body, deliveries, ts, self)
+
+    def get(
+        self,
+        block: bool = False,
+        timeout: Optional[float] = None,
+        poll_interval: float = 0.05,
+        lease_ms: Optional[int] = None,
+    ) -> Optional[Message]:
+        """Lease the next ready message. With ``block=True`` it polls until one
+        is available or ``timeout`` seconds elapse (``None`` = forever).
+        Returns ``None`` if empty/timed out."""
+        msg = self._wrap(self._core.dequeue(lease_ms))
+        if msg is not None or not block:
+            return msg
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            time.sleep(poll_interval)
+            msg = self._wrap(self._core.dequeue(lease_ms))
+            if msg is not None:
+                return msg
+            if deadline is not None and time.monotonic() >= deadline:
+                return None
+
+    def get_many(self, max_messages: int, lease_ms: Optional[int] = None) -> List[Message]:
+        """Lease up to ``max_messages`` ready messages at once."""
+        return [self._wrap(r) for r in self._core.dequeue_batch(max_messages, lease_ms)]
+
+    def consume(
+        self,
+        lease_ms: Optional[int] = None,
+        block: bool = True,
+        timeout: Optional[float] = None,
+        poll_interval: float = 0.05,
+        auto_ack: bool = True,
+    ) -> Iterator[Message]:
+        """Yield messages until the queue drains (or ``timeout`` elapses with
+        ``block=True``). Each message is auto-acked once the loop body finishes
+        cleanly, or nacked-for-redelivery if the body raises."""
+        while True:
+            msg = self.get(block=block, timeout=timeout, poll_interval=poll_interval, lease_ms=lease_ms)
+            if msg is None:
+                return
+            try:
+                yield msg
+            except GeneratorExit:
+                if not msg.settled:
+                    msg.nack(requeue=True)
+                raise
+            except Exception:
+                if not msg.settled:
+                    msg.nack(requeue=True)
+                raise
+            else:
+                if auto_ack and not msg.settled:
+                    msg.ack()
+
+    @contextmanager
+    def transaction(self, lease_ms: Optional[int] = None, requeue_on_error: bool = True):
+        """Context manager that leases one message, acks on success, and nacks
+        (requeue by default) on error. Yields ``None`` if the queue is empty."""
+        msg = self.get(block=False, lease_ms=lease_ms)
+        if msg is None:
+            yield None
+            return
+        try:
+            yield msg
+        except Exception:
+            if not msg.settled:
+                msg.nack(requeue=requeue_on_error)
+            raise
+        else:
+            if not msg.settled:
+                msg.ack()
+
+    # --- settling (also reachable directly via Message) ---
+
+    def ack(self, seq: int) -> bool:
+        return self._core.ack(seq)
+
+    def ack_many(self, seqs: List[int]) -> int:
+        return self._core.ack_batch(seqs)
+
+    def nack(self, seq: int, requeue: bool = True) -> bool:
+        return self._core.nack(seq, requeue)
+
+    def extend_lease(self, seq: int, lease_ms: Optional[int] = None) -> bool:
+        return self._core.extend_lease(seq, lease_ms)
+
+    # --- introspection / maintenance ---
+
+    def available(self) -> int:
+        """Messages ready for immediate delivery."""
+        return self._core.available()
+
+    def in_flight(self) -> int:
+        """Leased messages awaiting ack."""
+        return self._core.in_flight()
+
+    def __len__(self) -> int:
+        return self._core.available()
+
+    def stats(self) -> Dict[str, Any]:
+        return dict(self._core.stats())
+
+    def flush(self) -> None:
+        """Force an fsync of any buffered records."""
+        if not self._closed:
+            self._core.flush()
+
+    def compact(self) -> None:
+        """Rewrite the log, reclaiming space from acked messages."""
+        self._core.compact()
+
+    def purge(self) -> None:
+        """Delete every message and reset the log."""
+        self._core.purge()
+
+    def verify_integrity(self) -> Dict[str, Any]:
+        """CRC-verify the whole on-disk log. Returns ``{ok, records, bytes, message}``."""
+        return dict(self._core.verify_integrity())
+
+    def close(self) -> None:
+        if not self._closed:
+            try:
+                self._core.flush()
+            finally:
+                self._closed = True
+
+    def __enter__(self) -> "Queue":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()
+
+    def __repr__(self) -> str:
+        return f"Queue(path={self.path!r}, available={self.available()})"
+
+
+class Broker:
+    """Manages many named queues under a single root directory.
+
+        broker = Broker("./broker")
+        broker.queue("emails").put_json({"to": "a@b.c"})
+        broker.queue("thumbnails").put(b"...")
+    """
+
+    _SAFE = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.")
+
+    def __init__(self, root: str, **queue_defaults: Any):
+        self.root = root
+        self._defaults = queue_defaults
+        self._queues: Dict[str, Queue] = {}
+        self._lock = threading.RLock()
+        if root != ":memory:":
+            os.makedirs(root, exist_ok=True)
+
+    def _safe_name(self, name: str) -> str:
+        if not name or any(c not in self._SAFE for c in name) or name in (".", ".."):
+            raise ValueError(
+                f"invalid queue name {name!r}: use letters, digits, '-', '_', '.'"
+            )
+        return name
+
+    def queue(self, name: str, **opts: Any) -> Queue:
+        """Get (creating on first use) the named queue. Cached per name."""
+        name = self._safe_name(name)
+        with self._lock:
+            q = self._queues.get(name)
+            if q is None:
+                path = ":memory:" if self.root == ":memory:" else os.path.join(self.root, name)
+                merged = {**self._defaults, **opts}
+                q = Queue(path, **merged)
+                self._queues[name] = q
+            return q
+
+    def queue_names(self) -> List[str]:
+        """List on-disk queue names (directories containing a queue.log)."""
+        if self.root == ":memory:":
+            return list(self._queues.keys())
+        names = []
+        try:
+            for entry in os.scandir(self.root):
+                if entry.is_dir() and os.path.exists(os.path.join(entry.path, "queue.log")):
+                    names.append(entry.name)
+        except FileNotFoundError:
+            pass
+        return sorted(names)
+
+    def flush(self) -> None:
+        with self._lock:
+            for q in self._queues.values():
+                q.flush()
+
+    def close(self) -> None:
+        with self._lock:
+            for q in self._queues.values():
+                q.close()
+            self._queues.clear()
+
+    def __enter__(self) -> "Broker":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()
+
+    def __repr__(self) -> str:
+        return f"Broker(root={self.root!r}, queues={self.queue_names()})"