beaver-db 2.0rc2__py3-none-any.whl

beaver/channels.py

@@ -0,0 +1,249 @@
+import asyncio
+import json
+import time
+import weakref
+from typing import (
+    IO,
+    Any,
+    Iterator,
+    AsyncIterator,
+    Protocol,
+    runtime_checkable,
+    TYPE_CHECKING,
+    NamedTuple,
+)
+
+from pydantic import BaseModel
+
+from .manager import AsyncBeaverBase, atomic, emits
+
+if TYPE_CHECKING:
+    from .core import AsyncBeaverDB
+
+
+class ChannelMessage[T](BaseModel):
+    """A message received from a channel."""
+
+    channel: str
+    payload: T
+    timestamp: float
+
+
+@runtime_checkable
+class IBeaverChannel[T: BaseModel](Protocol):
+    """
+    The Synchronous Protocol exposed to the user via BeaverBridge.
+    """
+
+    def publish(self, payload: T) -> None: ...
+    def subscribe(self) -> Iterator[ChannelMessage]: ...
+    def history(self, limit: int = 100) -> list[ChannelMessage]: ...
+    def clear(self) -> None: ...
+    def count(self) -> int: ...
+
+
+class PubSubEngine:
+    """
+    A central engine that manages the background polling loop for ALL channels.
+    Attached to the AsyncBeaverDB instance.
+    """
+
+    def __init__(self, db: "AsyncBeaverDB"):
+        self.db = db
+        self._listeners: dict[str, list[asyncio.Queue]] = {}
+        self._running = False
+        self._task: asyncio.Task | None = None
+        self._last_poll_ts = time.time()
+
+    async def start(self):
+        """Starts the background polling loop."""
+        if self._running:
+            return
+        self._running = True
+        self._last_poll_ts = time.time()
+        self._task = asyncio.create_task(self._poll_loop())
+
+    async def stop(self):
+        """Stops the background polling loop."""
+        self._running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+
+    def subscribe(self, channel: str) -> asyncio.Queue[ChannelMessage]:
+        """Registers a new listener queue for a channel."""
+        queue = asyncio.Queue()
+        if channel not in self._listeners:
+            self._listeners[channel] = []
+        self._listeners[channel].append(queue)
+        return queue
+
+    def unsubscribe(self, channel: str, queue: asyncio.Queue):
+        """Unregisters a listener."""
+        if channel in self._listeners:
+            if queue in self._listeners[channel]:
+                self._listeners[channel].remove(queue)
+            if not self._listeners[channel]:
+                del self._listeners[channel]
+
+    async def _poll_loop(self):
+        """
+        Periodically checks the DB for new messages and dispatches them
+        to registered local queues.
+        """
+        while self._running:
+            try:
+                # 1. Fetch new messages globally
+                # We use a raw execute here to avoid locking the transaction logic
+                # for simple reads.
+                cursor = await self.db.connection.execute(
+                    """
+                    SELECT timestamp, channel_name, message_payload
+                    FROM __beaver_pubsub_log__
+                    WHERE timestamp > ?
+                    ORDER BY timestamp ASC
+                    """,
+                    (self._last_poll_ts,),
+                )
+
+                rows = await cursor.fetchall()
+                rows = list(rows)
+
+                if rows:
+                    # Update high-water mark
+                    self._last_poll_ts = rows[-1]["timestamp"]
+
+                    # 2. Dispatch to listeners
+                    for row in rows:
+                        channel = row["channel_name"]
+                        payload = row["message_payload"]
+                        msg = ChannelMessage(
+                            channel=channel, payload=payload, timestamp=row["timestamp"]
+                        )
+
+                        # Fan-out to all queues listening on this channel
+                        if channel in self._listeners:
+                            for q in self._listeners[channel]:
+                                q.put_nowait(msg)
+
+                # 3. Wait before next poll
+                # Adaptive sleep could be added here (e.g. sleep less if busy)
+                await asyncio.sleep(0.1)
+
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                # Log error or retry, don't crash the loop
+                await asyncio.sleep(1.0)
+
+
+class AsyncBeaverChannel[T: BaseModel](AsyncBeaverBase[T]):
+    """
+    A wrapper for a Pub/Sub channel.
+    Refactored for Async-First architecture (v2.0).
+    """
+
+    async def _get_engine(self) -> PubSubEngine:
+        """
+        Retrieves (or creates) the shared PubSubEngine on the DB instance.
+        """
+        # We perform a lazy attachment to the DB instance to avoid modifying core.py
+        # heavily. We store the engine on the DB instance dynamically.
+        if not hasattr(self._db, "_pubsub_engine"):
+            self._db._pubsub_engine = PubSubEngine(self._db)
+            await self._db._pubsub_engine.start()
+
+        return self._db._pubsub_engine
+
+    @emits("publish", payload=lambda payload, *args, **kwargs: dict(payload=payload))
+    @atomic
+    async def publish(self, payload: T):
+        """
+        Publishes a message to the channel.
+        """
+        # Ensure engine is running (in case we are the first publisher)
+        await self._get_engine()
+
+        # Serialize
+        data_str = self._serialize(payload)
+        ts = time.time()
+
+        # Monotonicity check (simple collision avoidance)
+        # In high-throughput, we rely on the DB to serialize inserts via lock
+        await self.connection.execute(
+            "INSERT INTO __beaver_pubsub_log__ (timestamp, channel_name, message_payload) VALUES (?, ?, ?)",
+            (ts, self._name, data_str),
+        )
+
+    async def subscribe(self) -> AsyncIterator[ChannelMessage[T]]:
+        """
+        Returns an async iterator that yields new messages as they arrive.
+        """
+        engine = await self._get_engine()
+        queue = engine.subscribe(self._name)
+
+        try:
+            while True:
+                # Wait for next message from the engine
+                msg = await queue.get()
+
+                # Deserialize message
+                yield ChannelMessage(
+                    channel=msg.channel,
+                    payload=self._deserialize(msg.payload),
+                    timestamp=msg.timestamp,
+                )
+        finally:
+            # Cleanup on break/cancel
+            engine.unsubscribe(self._name, queue)
+
+    async def history(self, limit: int = 100) -> list[ChannelMessage[T]]:
+        """
+        Retrieves the last N messages from the channel history.
+        """
+        cursor = await self.connection.execute(
+            """
+            SELECT timestamp, channel_name, message_payload
+            FROM __beaver_pubsub_log__
+            WHERE channel_name = ?
+            ORDER BY timestamp DESC
+            LIMIT ?
+            """,
+            (self._name, limit),
+        )
+
+        rows = await cursor.fetchall()
+        rows = list(rows)
+        results = []
+
+        # Rows are DESC, we usually want them returned chronologically
+        for row in reversed(rows):
+            payload = self._deserialize(row["message_payload"])
+            results.append(
+                ChannelMessage(
+                    channel=self._name, payload=payload, timestamp=row["timestamp"]
+                )
+            )
+
+        return results
+
+    @emits("clear", payload=lambda *args, **kwargs: dict())
+    @atomic
+    async def clear(self):
+        """Clears the history for this channel."""
+        await self.connection.execute(
+            "DELETE FROM __beaver_pubsub_log__ WHERE channel_name = ?", (self._name,)
+        )
+
+    async def count(self) -> int:
+        """Returns the total number of messages in history."""
+        cursor = await self.connection.execute(
+            "SELECT COUNT(*) FROM __beaver_pubsub_log__ WHERE channel_name = ?",
+            (self._name,),
+        )
+        row = await cursor.fetchone()
+        return row[0] if row else 0

beaver/cli/__init__.py

@@ -0,0 +1,133 @@
+import typer
+import rich
+import rich.table
+from typing_extensions import Annotated
+
+import beaver
+from beaver import BeaverDB
+
+# Import the command group from the new file
+from beaver.cli import dicts as dicts_cli
+from beaver.cli import lists as lists_cli
+from beaver.cli import queues as queues_cli
+from beaver.cli import blobs as blobs_cli
+from beaver.cli import locks as locks_cli
+from beaver.cli import logs as logs_cli
+from beaver.cli import channels as channels_cli
+from beaver.cli import collections as collections_cli
+
+# --- Main App ---
+app = typer.Typer()
+
+# Register the command group
+app.add_typer(dicts_cli.app)
+app.add_typer(lists_cli.app)
+app.add_typer(queues_cli.app)
+app.add_typer(blobs_cli.app)
+app.add_typer(locks_cli.app)
+app.add_typer(logs_cli.app)
+app.add_typer(channels_cli.app)
+app.add_typer(collections_cli.app)
+
+
+def version_callback(value: bool):
+    if value:
+        print(beaver.__version__)
+        raise typer.Exit()
+
+
+@app.callback()
+def main(
+    ctx: typer.Context,
+    database: Annotated[
+        str, typer.Option(help="The path to the BeaverDB database file.")
+    ] = "beaver.db",
+    version: Annotated[
+        bool,
+        typer.Option(
+            "--version",
+            callback=version_callback,
+            is_eager=True,
+            help="Show the version and exit.",
+        ),
+    ] = False,
+):
+    """
+    BeaverDB command-line interface.
+    """
+    try:
+        # Store the db instance in the context for all subcommands
+        ctx.obj = {"db": BeaverDB(database)}
+    except Exception as e:
+        rich.print(f"[bold red]Error opening database:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+# --- Serve Command ---
+@app.command()
+def serve(
+    database: Annotated[
+        str, typer.Option(help="The path to the BeaverDB database file.")
+    ] = "beaver.db",
+    host: Annotated[
+        str, typer.Option(help="The host to bind the server to.")
+    ] = "127.0.0.1",
+    port: Annotated[int, typer.Option(help="The port to run the server on.")] = 8000,
+):
+    """Starts a REST API server for the BeaverDB database."""
+    try:
+        from beaver import server
+    except ImportError:
+        rich.print(
+            "[red]Error:[/] To use the serve command, please install the server dependencies:\n"
+            'pip install "beaver-db[server]"'
+        )
+        raise typer.Exit(code=1)
+    server.serve(database, host=host, port=port)
+
+
+@app.command()
+def info(ctx: typer.Context):
+    """
+    Displays a summary of all data structures in the database.
+    """
+    db: BeaverDB = ctx.obj["db"]
+    rich.print(f"[bold]Database Summary: {db._db_path}[/bold]")
+
+    table = rich.table.Table(title="BeaverDB Contents")
+    table.add_column("Type", style="cyan")
+    table.add_column("Name", style="bold")
+    table.add_column("Count", style="magenta", justify="right")
+
+    try:
+        # Dictionaries
+        for name in db.dicts:
+            table.add_row("Dict", name, str(len(db.dict(name))))
+        # Lists
+        for name in db.lists:
+            table.add_row("List", name, str(len(db.list(name))))
+        # Queues
+        for name in db.queues:
+            table.add_row("Queue", name, str(len(db.queue(name))))
+        # Collections
+        for name in db.collections:
+            table.add_row("Collection", name, str(len(db.collection(name))))
+        # Blob Stores
+        for name in db.blobs:
+            table.add_row("Blob Store", name, str(len(db.blob(name))))
+        # Logs
+        for name in db.logs:
+            table.add_row("Log", name, "N/A (len not supported)")
+        # Active Locks
+        for name in db.locks:
+            table.add_row("Active Lock", name, "1")
+
+        rich.print(table)
+
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+if __name__ == "__main__":
+    app()

beaver/cli/blobs.py

@@ -0,0 +1,225 @@
+import json
+from pathlib import Path
+import typer
+import rich
+import rich.table
+from typing_extensions import Annotated
+from typing import Optional
+
+from beaver import BeaverDB
+
+app = typer.Typer(
+    name="blob",
+    help="Interact with blob stores. (e.g., beaver blob assets put my-file.png /path/to/file.png)",
+)
+
+
+def _get_db(ctx: typer.Context) -> BeaverDB:
+    """Helper to get the DB instance from the main context."""
+    return ctx.find_object(dict)["db"]
+
+
+def _parse_metadata(metadata_str: Optional[str]) -> Optional[dict]:
+    """Parses the metadata string as JSON."""
+    if metadata_str is None:
+        return None
+    if not (metadata_str.startswith("{") or metadata_str.startswith("[")):
+        raise typer.BadParameter(
+            'Metadata must be valid JSON (e.g., \'{"key":"value"}\')'
+        )
+    try:
+        return json.loads(metadata_str)
+    except json.JSONDecodeError:
+        raise typer.BadParameter("Invalid JSON format for metadata.")
+
+
+@app.callback(invoke_without_command=True)
+def blob_main(
+    ctx: typer.Context,
+    name: Annotated[
+        Optional[str],
+        typer.Argument(help="The name of the blob store to interact with."),
+    ] = None,
+):
+    """
+    Manage binary blob stores.
+
+    If no name is provided, lists all available blob stores.
+    """
+    db = _get_db(ctx)
+
+    if name is None:
+        # No name given, so list all blob stores
+        rich.print("[bold]Available Blob Stores:[/bold]")
+        try:
+            blob_names = db.blobs
+            if not blob_names:
+                rich.print("  (No blob stores found)")
+            else:
+                for blob_name in blob_names:
+                    rich.print(f"  • {blob_name}")
+            rich.print(
+                "\n[bold]Usage:[/bold] beaver blob [bold]<NAME>[/bold] [COMMAND]"
+            )
+            return
+        except Exception as e:
+            rich.print(f"[bold red]Error querying blob stores:[/] {e}")
+            raise typer.Exit(code=1)
+
+    # A name was provided, store it in the context for subcommands
+    ctx.obj = {"name": name, "db": db}
+
+    if ctx.invoked_subcommand is None:
+        # A name was given, but no command
+        try:
+            count = len(db.blob(name))
+            rich.print(f"Blob Store '[bold]{name}[/bold]' contains {count} items.")
+            rich.print("\n[bold]Commands:[/bold]")
+            rich.print("  put, get, del, list, dump")
+            rich.print(
+                f"\nRun [bold]beaver blob {name} --help[/bold] for command-specific options."
+            )
+        except Exception as e:
+            rich.print(f"[bold red]Error:[/] {e}")
+            raise typer.Exit(code=1)
+        raise typer.Exit()
+
+
+@app.command()
+def put(
+    ctx: typer.Context,
+    key: Annotated[str, typer.Argument(help="The unique key to store the blob under.")],
+    file_path: Annotated[
+        Path,
+        typer.Argument(
+            exists=True,
+            dir_okay=False,
+            readable=True,
+            help="The path to the file to upload.",
+        ),
+    ],
+    metadata: Annotated[
+        Optional[str], typer.Option(help="Optional metadata as a JSON string.")
+    ] = None,
+):
+    """
+    Put (upload) a file into the blob store.
+    """
+    db = ctx.obj["db"]
+    name = ctx.obj["name"]
+    try:
+        parsed_metadata = _parse_metadata(metadata)
+
+        with open(file_path, "rb") as f:
+            file_bytes = f.read()
+
+        db.blob(name).put(key, file_bytes, metadata=parsed_metadata)
+        rich.print(
+            f"[green]Success:[/] File '{file_path}' stored as key '{key}' in blob store '{name}'."
+        )
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command()
+def get(
+    ctx: typer.Context,
+    key: Annotated[str, typer.Argument(help="The key of the blob to retrieve.")],
+    output_path: Annotated[
+        Path,
+        typer.Argument(
+            exists=True,
+            dir_okay=False,
+            readable=True,
+            help="The path to the file to upload.",
+        ),
+    ],
+):
+    """
+    Get (download) a file from the blob store.
+    """
+    db = ctx.obj["db"]
+    name = ctx.obj["name"]
+    try:
+        blob = db.blob(name).get(key)
+        if blob is None:
+            rich.print(f"[bold red]Error:[/] Key not found: '{key}'")
+            raise typer.Exit(code=1)
+
+        with open(output_path, "wb") as f:
+            f.write(blob.data)
+
+        rich.print(f"[green]Success:[/] Blob '{key}' saved to '{output_path}'.")
+        if blob.metadata:
+            rich.print("[bold]Metadata:[/bold]")
+            if isinstance(blob.metadata, (dict, list)):
+                rich.print_json(data=blob.metadata)
+            else:
+                rich.print(str(blob.metadata))
+
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command(name="del")
+def delete(
+    ctx: typer.Context,
+    key: Annotated[str, typer.Argument(help="The key of the blob to delete.")],
+):
+    """
+    Delete a blob from the store.
+    """
+    db = ctx.obj["db"]
+    name = ctx.obj["name"]
+    try:
+        db.blob(name).delete(key)
+        rich.print(f"[green]Success:[/] Blob '{key}' deleted from store '{name}'.")
+    except KeyError:
+        rich.print(f"[bold red]Error:[/] Key not found: '{key}'")
+        raise typer.Exit(code=1)
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command(name="list")
+def list_keys(ctx: typer.Context):
+    """
+    List all keys in the blob store.
+    """
+    db = ctx.obj["db"]
+    name = ctx.obj["name"]
+    try:
+        # The __iter__ for BlobManager yields keys
+        all_keys = list(db.blob(name))
+        if not all_keys:
+            rich.print(f"Blob store '{name}' is empty.")
+            return
+
+        table = rich.table.Table(title=f"Keys in Blob Store: [bold]{name}[/bold]")
+        table.add_column("Key")
+
+        for key in all_keys:
+            table.add_row(key)
+        rich.print(table)
+
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command()
+def dump(ctx: typer.Context):
+    """
+    Dump the entire blob store as JSON (with data as base64).
+    """
+    db = ctx.obj["db"]
+    name = ctx.obj["name"]
+    try:
+        dump_data = db.blob(name).dump()
+        rich.print_json(data=dump_data)
+    except Exception as e:
+        rich.print(f"[bold red]Error:[/] {e}")
+        raise typer.Exit(code=1)