pydocket 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

docket/cli.py

@@ -1,15 +1,21 @@
 import asyncio
 import enum
+import importlib
 import logging
+import os
 import socket
 import sys
-from datetime import timedelta
-from typing import Annotated
+from datetime import datetime, timedelta, timezone
+from functools import partial
+from typing import Annotated, Any, Collection
 
 import typer
+from rich.console import Console
+from rich.table import Table
 
 from . import __version__, tasks
-from .docket import Docket
+from .docket import Docket, DocketSnapshot, WorkerInfo
+from .execution import Operator
 from .worker import Worker
 
 app: typer.Typer = typer.Typer(
@@ -33,6 +39,14 @@ class LogFormat(enum.StrEnum):
     JSON = "json"
 
 
+def local_time(when: datetime) -> str:
+    return when.astimezone().strftime("%Y-%m-%d %H:%M:%S %z")
+
+
+def default_worker_name() -> str:
+    return f"{socket.gethostname()}#{os.getpid()}"
+
+
 def duration(duration_str: str | timedelta) -> timedelta:
     """
     Parse a duration string into a timedelta.
@@ -100,6 +114,41 @@ def set_logging_level(level: LogLevel) -> None:
     logging.getLogger().setLevel(level)
 
 
+def handle_strike_wildcard(value: str) -> str | None:
+    if value in ("", "*"):
+        return None
+    return value
+
+
+def interpret_python_value(value: str | None) -> Any:
+    if value is None:
+        return None
+
+    type, _, value = value.rpartition(":")
+    if not type:
+        # without a type hint, we assume the value is a string
+        return value
+
+    module_name, _, member_name = type.rpartition(".")
+    module = importlib.import_module(module_name or "builtins")
+    member = getattr(module, member_name)
+
+    # special cases for common useful types
+    if member is timedelta:
+        return timedelta(seconds=int(value))
+    elif member is bool:
+        return value.lower() == "true"
+    else:
+        return member(value)
+
+
+@app.command(
+    help="Print the version of docket",
+)
+def version() -> None:
+    print(__version__)
+
+
 @app.command(
     help="Start a worker to process tasks",
 )
@@ -136,7 +185,7 @@ def worker(
             help="The name of the worker",
             envvar="DOCKET_WORKER_NAME",
         ),
-    ] = socket.gethostname(),
+    ] = default_worker_name(),
     logging_level: Annotated[
         LogLevel,
         typer.Option(
@@ -153,11 +202,11 @@ def worker(
             callback=set_logging_format,
         ),
     ] = LogFormat.RICH if sys.stdout.isatty() else LogFormat.PLAIN,
-    prefetch_count: Annotated[
+    concurrency: Annotated[
         int,
         typer.Option(
-            help="The number of tasks to request from the docket at a time",
-            envvar="DOCKET_WORKER_PREFETCH_COUNT",
+            help="The maximum number of tasks to process concurrently",
+            envvar="DOCKET_WORKER_CONCURRENCY",
         ),
     ] = 10,
     redelivery_timeout: Annotated[
@@ -192,7 +241,7 @@ def worker(
             docket_name=docket_,
             url=url,
             name=name,
-            prefetch_count=prefetch_count,
+            concurrency=concurrency,
             redelivery_timeout=redelivery_timeout,
             reconnection_delay=reconnection_delay,
             until_finished=until_finished,
@@ -201,7 +250,139 @@ def worker(
     )
 
 
-@app.command(help="Adds a trace task to the Docket")
+@app.command(help="Strikes a task or parameters from the docket")
+def strike(
+    function: Annotated[
+        str,
+        typer.Argument(
+            help="The function to strike",
+            callback=handle_strike_wildcard,
+        ),
+    ] = "*",
+    parameter: Annotated[
+        str,
+        typer.Argument(
+            help="The parameter to strike",
+            callback=handle_strike_wildcard,
+        ),
+    ] = "*",
+    operator: Annotated[
+        Operator,
+        typer.Argument(
+            help="The operator to compare the value against",
+        ),
+    ] = Operator.EQUAL,
+    value: Annotated[
+        str | None,
+        typer.Argument(
+            help="The value to strike from the docket",
+        ),
+    ] = None,
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+) -> None:
+    if not function and not parameter:
+        raise typer.BadParameter(
+            message="Must provide either a function and/or a parameter",
+        )
+
+    value_ = interpret_python_value(value)
+    if parameter:
+        function_name = f"{function or '(all tasks)'}"
+        print(f"Striking {function_name} {parameter} {operator} {value_!r}")
+    else:
+        print(f"Striking {function}")
+
+    async def run() -> None:
+        async with Docket(name=docket_, url=url) as docket:
+            await docket.strike(function, parameter, operator, value_)
+
+    asyncio.run(run())
+
+
+@app.command(help="Restores a task or parameters to the Docket")
+def restore(
+    function: Annotated[
+        str,
+        typer.Argument(
+            help="The function to restore",
+            callback=handle_strike_wildcard,
+        ),
+    ] = "*",
+    parameter: Annotated[
+        str,
+        typer.Argument(
+            help="The parameter to restore",
+            callback=handle_strike_wildcard,
+        ),
+    ] = "*",
+    operator: Annotated[
+        Operator,
+        typer.Argument(
+            help="The operator to compare the value against",
+        ),
+    ] = Operator.EQUAL,
+    value: Annotated[
+        str | None,
+        typer.Argument(
+            help="The value to restore to the docket",
+        ),
+    ] = None,
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+) -> None:
+    if not function and not parameter:
+        raise typer.BadParameter(
+            message="Must provide either a function and/or a parameter",
+        )
+
+    value_ = interpret_python_value(value)
+    if parameter:
+        function_name = f"{function or '(all tasks)'}"
+        print(f"Striking {function_name} {parameter} {operator} {value_!r}")
+    else:
+        print(f"Restoring {function}")
+
+    async def run() -> None:
+        async with Docket(name=docket_, url=url) as docket:
+            await docket.restore(function, parameter, operator, value_)
+
+    asyncio.run(run())
+
+
+tasks_app: typer.Typer = typer.Typer(
+    help="Run docket's built-in tasks", no_args_is_help=True
+)
+app.add_typer(tasks_app, name="tasks")
+
+
+@tasks_app.command(help="Adds a trace task to the Docket")
 def trace(
     docket_: Annotated[
         str,
@@ -224,21 +405,104 @@ def trace(
             help="The message to print",
         ),
     ] = "Howdy!",
-    error: Annotated[
-        bool,
+    delay: Annotated[
+        timedelta,
         typer.Option(
-            "--error",
-            help="Intentionally raise an error",
+            parser=duration,
+            help="The delay before the task is added to the docket",
         ),
-    ] = False,
+    ] = timedelta(seconds=0),
+) -> None:
+    async def run() -> None:
+        async with Docket(name=docket_, url=url) as docket:
+            when = datetime.now(timezone.utc) + delay
+            execution = await docket.add(tasks.trace, when)(message)
+            print(
+                f"Added {execution.function.__name__} task {execution.key!r} to "
+                f"the docket {docket.name!r}"
+            )
+
+    asyncio.run(run())
+
+
+@tasks_app.command(help="Adds a fail task to the Docket")
+def fail(
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+    message: Annotated[
+        str,
+        typer.Argument(
+            help="The message to print",
+        ),
+    ] = "Howdy!",
+    delay: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help="The delay before the task is added to the docket",
+        ),
+    ] = timedelta(seconds=0),
 ) -> None:
     async def run() -> None:
         async with Docket(name=docket_, url=url) as docket:
-            if error:
-                execution = await docket.add(tasks.fail)(message)
-            else:
-                execution = await docket.add(tasks.trace)(message)
+            when = datetime.now(timezone.utc) + delay
+            execution = await docket.add(tasks.fail, when)(message)
+            print(
+                f"Added {execution.function.__name__} task {execution.key!r} to "
+                f"the docket {docket.name!r}"
+            )
+
+    asyncio.run(run())
+
 
+@tasks_app.command(help="Adds a sleep task to the Docket")
+def sleep(
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+    seconds: Annotated[
+        float,
+        typer.Argument(
+            help="The number of seconds to sleep",
+        ),
+    ] = 1,
+    delay: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help="The delay before the task is added to the docket",
+        ),
+    ] = timedelta(seconds=0),
+) -> None:
+    async def run() -> None:
+        async with Docket(name=docket_, url=url) as docket:
+            when = datetime.now(timezone.utc) + delay
+            execution = await docket.add(tasks.sleep, when)(seconds)
             print(
                 f"Added {execution.function.__name__} task {execution.key!r} to "
                 f"the docket {docket.name!r}"
@@ -247,8 +511,179 @@ def trace(
     asyncio.run(run())
 
 
-@app.command(
-    help="Print the version of Docket",
+def relative_time(now: datetime, when: datetime) -> str:
+    delta = now - when
+    if delta < -timedelta(minutes=30):
+        return f"at {local_time(when)}"
+    elif delta < timedelta(0):
+        return f"in {-delta}"
+    elif delta < timedelta(minutes=30):
+        return f"{delta} ago"
+    else:
+        return f"at {local_time(when)}"
+
+
+@app.command(help="Shows a snapshot of what's on the docket right now")
+def snapshot(
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+) -> None:
+    async def run() -> DocketSnapshot:
+        async with Docket(name=docket_, url=url) as docket:
+            return await docket.snapshot()
+
+    snapshot = asyncio.run(run())
+
+    relative = partial(relative_time, snapshot.taken)
+
+    console = Console()
+
+    summary_lines = [
+        f"Docket: {docket_!r}",
+        f"as of {local_time(snapshot.taken)}",
+        (
+            f"{len(snapshot.workers)} workers, "
+            f"{len(snapshot.running)}/{snapshot.total_tasks} running"
+        ),
+    ]
+    table = Table(title="\n".join(summary_lines))
+    table.add_column("When", style="green")
+    table.add_column("Function", style="cyan")
+    table.add_column("Key", style="cyan")
+    table.add_column("Worker", style="yellow")
+    table.add_column("Started", style="green")
+
+    for execution in snapshot.running:
+        table.add_row(
+            relative(execution.when),
+            execution.function.__name__,
+            execution.key,
+            execution.worker,
+            relative(execution.started),
+        )
+
+    for execution in snapshot.future:
+        table.add_row(
+            relative(execution.when),
+            execution.function.__name__,
+            execution.key,
+            "",
+            "",
+        )
+
+    console.print(table)
+
+
+workers_app: typer.Typer = typer.Typer(
+    help="Look at the workers on a docket", no_args_is_help=True
 )
-def version() -> None:
-    print(__version__)
+app.add_typer(workers_app, name="workers")
+
+
+def print_workers(
+    docket_name: str,
+    workers: Collection[WorkerInfo],
+    highlight_task: str | None = None,
+) -> None:
+    sorted_workers = sorted(workers, key=lambda w: w.last_seen, reverse=True)
+
+    table = Table(title=f"Workers in Docket: {docket_name}")
+
+    table.add_column("Name", style="cyan")
+    table.add_column("Last Seen", style="green")
+    table.add_column("Tasks", style="yellow")
+
+    now = datetime.now(timezone.utc)
+
+    for worker in sorted_workers:
+        time_ago = now - worker.last_seen
+
+        tasks = [
+            f"[bold]{task}[/bold]" if task == highlight_task else task
+            for task in sorted(worker.tasks)
+        ]
+
+        table.add_row(
+            worker.name,
+            f"{time_ago} ago",
+            "\n".join(tasks) if tasks else "(none)",
+        )
+
+    console = Console()
+    console.print(table)
+
+
+@workers_app.command(name="ls", help="List all workers on the docket")
+def list_workers(
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+) -> None:
+    async def run() -> Collection[WorkerInfo]:
+        async with Docket(name=docket_, url=url) as docket:
+            return await docket.workers()
+
+    workers = asyncio.run(run())
+
+    print_workers(docket_, workers)
+
+
+@workers_app.command(
+    name="for-task",
+    help="List the workers on the docket that can process a certain task",
+)
+def workers_for_task(
+    task: Annotated[
+        str,
+        typer.Argument(
+            help="The name of the task",
+        ),
+    ],
+    docket_: Annotated[
+        str,
+        typer.Option(
+            "--docket",
+            help="The name of the docket",
+            envvar="DOCKET_NAME",
+        ),
+    ] = "docket",
+    url: Annotated[
+        str,
+        typer.Option(
+            help="The URL of the Redis server",
+            envvar="DOCKET_URL",
+        ),
+    ] = "redis://localhost:6379/0",
+) -> None:
+    async def run() -> Collection[WorkerInfo]:
+        async with Docket(name=docket_, url=url) as docket:
+            return await docket.task_workers(task)
+
+    workers = asyncio.run(run())
+
+    print_workers(docket_, workers, highlight_task=task)

docket/dependencies.py

@@ -59,7 +59,7 @@ def TaskKey() -> str:
 class _TaskLogger(Dependency):
     def __call__(
         self, docket: Docket, worker: Worker, execution: Execution
-    ) -> logging.LoggerAdapter:
+    ) -> logging.LoggerAdapter[logging.Logger]:
         logger = logging.getLogger(f"docket.task.{execution.function.__name__}")
 
         extra = {