beaver-db 2.0rc2__py3-none-any.whl

beaver/queues.py

@@ -0,0 +1,215 @@
+import asyncio
+import json
+import time
+from datetime import datetime, timezone
+from typing import (
+    IO,
+    Iterator,
+    Literal,
+    NamedTuple,
+    overload,
+    Protocol,
+    runtime_checkable,
+    TYPE_CHECKING,
+)
+
+from pydantic import BaseModel
+
+from .manager import AsyncBeaverBase, atomic, emits
+
+if TYPE_CHECKING:
+    from .core import AsyncBeaverDB
+
+
+class QueueItem[T](NamedTuple):
+    """A data class representing a single item retrieved from the queue."""
+
+    priority: float
+    timestamp: float
+    data: T
+
+
+@runtime_checkable
+class IBeaverQueue[T: BaseModel](Protocol):
+    """
+    The Synchronous Protocol exposed to the user via BeaverBridge.
+    """
+
+    def put(self, data: T, priority: float) -> None: ...
+
+    def peek(self) -> QueueItem[T] | None: ...
+
+    # Overloads for get
+    @overload
+    def get(
+        self, block: Literal[True] = True, timeout: float | None = None
+    ) -> QueueItem[T]: ...
+    @overload
+    def get(self, block: Literal[False]) -> QueueItem[T]: ...
+
+    def clear(self) -> None: ...
+    def count(self) -> int: ...
+    def dump(self, fp: IO[str] | None = None) -> dict | None: ...
+
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[QueueItem[T]]: ...
+    def __bool__(self) -> bool: ...
+
+
+class AsyncBeaverQueue[T: BaseModel](AsyncBeaverBase[T]):
+    """
+    A wrapper providing a Pythonic interface to a persistent, multi-process
+    producer-consumer priority queue.
+    Refactored for Async-First architecture (v2.0).
+    """
+
+    @emits("put", payload=lambda *args, **kwargs: dict())
+    @atomic
+    async def put(self, data: T, priority: float):
+        """
+        Adds an item to the queue with a specific priority.
+        """
+        await self.connection.execute(
+            "INSERT INTO __beaver_priority_queues__ (queue_name, priority, timestamp, data) VALUES (?, ?, ?, ?)",
+            (self._name, priority, time.time(), self._serialize(data)),
+        )
+
+    async def _get_item_atomically(self, pop: bool = True) -> QueueItem[T] | None:
+        """
+        Performs a single, atomic attempt to retrieve and remove the
+        highest-priority item from the queue.
+        """
+        # We need a transaction to ensure we don't peek/get an item that someone else steals
+        # Note: If called from peek/get, they might handle locking via @atomic or logic.
+        # Since this helper is used inside @atomic methods, we just need execution.
+
+        cursor = await self.connection.execute(
+            """
+            SELECT rowid, priority, timestamp, data
+            FROM __beaver_priority_queues__
+            WHERE queue_name = ?
+            ORDER BY priority ASC, timestamp ASC
+            LIMIT 1
+            """,
+            (self._name,),
+        )
+        result = await cursor.fetchone()
+
+        if result is None:
+            return None
+
+        rowid, priority, timestamp, data = result
+
+        if pop:
+            await self.connection.execute(
+                "DELETE FROM __beaver_priority_queues__ WHERE rowid = ?", (rowid,)
+            )
+
+        return QueueItem(
+            priority=priority, timestamp=timestamp, data=self._deserialize(data)
+        )
+
+    @atomic
+    async def peek(self) -> QueueItem[T] | None:
+        """
+        Retrieves the first item of the queue without removing it.
+        """
+        return await self._get_item_atomically(pop=False)
+
+    @atomic
+    async def _try_pop_atomic(self) -> QueueItem[T] | None:
+        """Helper to check and pop one item under lock."""
+        return await self._get_item_atomically(pop=True)
+
+    async def _get_loop_impl(self, block: bool, timeout: float | None) -> QueueItem[T]:
+        """
+        The polling loop. It acquires the lock only briefly during the check,
+        allowing producers to interleave 'put' operations while we wait.
+        """
+        # 1. Non-blocking fast path
+        if not block:
+            item = await self._try_pop_atomic()
+            if item is None:
+                raise IndexError("get from an empty queue.")
+            return item
+
+        # 2. Blocking loop
+        start_time = time.time()
+        while True:
+            item = await self._try_pop_atomic()
+
+            if item is not None:
+                return item
+
+            if timeout is not None and (time.time() - start_time) > timeout:
+                raise TimeoutError("Timeout expired while waiting for an item.")
+
+            # Yield control to allow producers to acquire the lock and put items
+            await asyncio.sleep(0.1)
+
+    # We override the public get to use the loop implementation
+    # NOTE: We do NOT decorate this with @atomic, as it manages its own locking scope.
+    @emits("get", payload=lambda *args, **kwargs: dict())
+    async def get(
+        self, block: bool = True, timeout: float | None = None
+    ) -> QueueItem[T]:
+        return await self._get_loop_impl(block, timeout)
+
+    async def count(self) -> int:
+        cursor = await self.connection.execute(
+            "SELECT COUNT(*) FROM __beaver_priority_queues__ WHERE queue_name = ?",
+            (self._name,),
+        )
+        count = await cursor.fetchone()
+        return count[0] if count else 0
+
+    async def clear(self):
+        await self.connection.execute(
+            "DELETE FROM __beaver_priority_queues__ WHERE queue_name = ?",
+            (self._name,),
+        )
+
+    # --- Iterators ---
+
+    async def __aiter__(self):
+        cursor = await self.connection.execute(
+            """
+            SELECT priority, timestamp, data
+            FROM __beaver_priority_queues__
+            WHERE queue_name = ?
+            ORDER BY priority ASC, timestamp ASC
+            """,
+            (self._name,),
+        )
+        async for row in cursor:
+            yield QueueItem(
+                priority=row["priority"],
+                timestamp=row["timestamp"],
+                data=self._deserialize(row["data"]),
+            )
+
+    async def dump(self, fp: IO[str] | None = None) -> dict | None:
+        items_list = []
+        async for item in self:
+            data = item.data
+            if self._model and isinstance(data, BaseModel):
+                data = json.loads(data.model_dump_json())
+
+            items_list.append(
+                {"priority": item.priority, "timestamp": item.timestamp, "data": data}
+            )
+
+        metadata = {
+            "type": "Queue",
+            "name": self._name,
+            "count": len(items_list),
+            "dump_date": datetime.now(timezone.utc).isoformat(),
+        }
+
+        dump_obj = {"metadata": metadata, "items": items_list}
+
+        if fp:
+            json.dump(dump_obj, fp, indent=2)
+            return None
+
+        return dump_obj

beaver/security.py

@@ -0,0 +1,144 @@
+import base64
+import os
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+try:
+    from cryptography.fernet import Fernet, InvalidToken
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+    from cryptography.hazmat.backends import default_backend
+except ImportError:
+    # We define dummy classes or raise errors if accessed without dependencies
+    Fernet = None  # type: ignore
+
+    def require_security():
+        raise ImportError(
+            'The "security" extra is required to use this feature. '
+            'Please install it with: pip install "beaver-db[security]"'
+        )
+
+else:
+
+    def require_security():
+        pass
+
+
+class Cipher:
+    """
+    Handles symmetric encryption using Fernet (AES-128 CBC + HMAC-SHA256).
+    Used by Encrypted Dictionaries to secure values at rest.
+    """
+
+    def __init__(self, secret: str, salt: bytes | None = None):
+        require_security()
+        if not salt:
+            salt = os.urandom(16)
+        self.salt = salt
+        self.key = self._derive_key(secret, salt)
+        self.fernet = Fernet(self.key)
+
+    @staticmethod
+    def _derive_key(secret: str, salt: bytes) -> bytes:
+        """Derives a 32-byte, url-safe base64 key from the secret using PBKDF2."""
+        kdf = PBKDF2HMAC(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=salt,
+            iterations=100_000,
+            backend=default_backend(),
+        )
+        return base64.urlsafe_b64encode(kdf.derive(secret.encode()))
+
+    def encrypt(self, data: bytes) -> bytes:
+        """Encrypts bytes."""
+        return self.fernet.encrypt(data)
+
+    def decrypt(self, token: bytes) -> bytes:
+        """Decrypts bytes. Raises ValueError if the key is incorrect."""
+        try:
+            return self.fernet.decrypt(token)
+        except InvalidToken:
+            raise ValueError("Invalid secret key or corrupted data.")
+
+
+class Secret(BaseModel):
+    """
+    A value type for secure, one-way password hashing.
+    Stores only the hash and salt as base64 strings, ensuring JSON compatibility.
+    """
+
+    hash: str
+    salt: str
+
+    def __init__(self, value: str | None = None, **kwargs):
+        """
+        Can be initialized in two ways:
+        1. New Secret: Secret("my-password") -> Hashes and stores it as b64 strings.
+        2. Loading:    Secret(hash="...", salt="...") -> Reconstructs it.
+        """
+        # Ensure security deps are present
+        if value is not None or (not kwargs.get("hash")):
+            require_security()
+
+        if value is not None:
+            # 1. User provided a plain-text password to hash
+            salt_bytes = os.urandom(16)
+            kdf = PBKDF2HMAC(
+                algorithm=hashes.SHA256(),
+                length=32,
+                salt=salt_bytes,
+                iterations=100_000,
+                backend=default_backend(),
+            )
+            hashed_bytes = kdf.derive(value.encode())
+
+            # Store as base64 strings
+            super().__init__(
+                hash=base64.b64encode(hashed_bytes).decode("utf-8"),
+                salt=base64.b64encode(salt_bytes).decode("utf-8"),
+            )
+        else:
+            # 2. Pydantic is loading from arguments (hash/salt)
+            super().__init__(**kwargs)
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        Checks equality.
+        - If comparing to another Secret, checks hash equality.
+        - If comparing to a string, checks if the string hashes to the stored value.
+        """
+        if isinstance(other, Secret):
+            return self.hash == other.hash and self.salt == other.salt
+
+        if isinstance(other, str):
+            require_security()
+
+            # Decode stored strings back to bytes
+            try:
+                salt_bytes = base64.b64decode(self.salt)
+                hash_bytes = base64.b64decode(self.hash)
+            except (ValueError, TypeError):
+                return False
+
+            kdf = PBKDF2HMAC(
+                algorithm=hashes.SHA256(),
+                length=32,
+                salt=salt_bytes,
+                iterations=100_000,
+                backend=default_backend(),
+            )
+            try:
+                kdf.verify(other.encode(), hash_bytes)
+                return True
+            except Exception:
+                return False
+
+        return False
+
+    def __repr__(self) -> str:
+        return f"Secret(hash={self.hash[:8]}...)"
+
+    def __str__(self) -> str:
+        return "********"