Merge pull request #1 from cofin/config-updates

fix: config and CLI updates

Author: cofin

README.md

@@ -17,9 +17,21 @@ Here is a basic application that demonstrates how to use the plugin.
 from __future__ import annotations
 
 from litestar import Litestar
-from litestar_saq import SaqPlugin, SaqConfig
 
-saq = SaqPlugin(config=SaqConfig())
+from litestar_saq import QueueConfig, SAQConfig, SAQPlugin
+
+saq = SAQPlugin(config=SAQConfig(redis_url="redis://localhost:6397/0", queue_configs=[QueueConfig(name="samples")]))
 app = Litestar(plugins=[saq])
 
+
+```
+
+You can start a background worker with the following command now:
+
+```shell
+litestar --app-dir=examples/ --app basic:app tasks run-worker
+Using Litestar app from env: 'basic:app'
+Starting SAQ Workers ──────────────────────────────────────────────────────────────────
+INFO - 2023-10-04 17:39:03,255 - saq - worker - Worker starting: Queue<redis=Redis<ConnectionPool<Connection<host=localhost,port=6397,db=0>>>, name='samples'>
+INFO - 2023-10-04 17:39:06,545 - saq - worker - Worker shutting down
 ```

examples/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from litestar import Controller, Litestar, get
+
+from litestar_saq import QueueConfig, SAQConfig, SAQPlugin
+
+if TYPE_CHECKING:
+    from saq.types import QueueInfo
+
+    from litestar_saq.base import Queue
+
+
+class SampleController(Controller):
+    @get(path="/samples")
+    async def samples_queue_info(self, task_queues: dict[str, Queue]) -> QueueInfo:
+        """Check database available and returns app config info."""
+        queue = task_queues.get("samples")
+        return await queue.info()  # type: ignore[union-attr]
+
+
+saq = SAQPlugin(config=SAQConfig(redis_url="redis://localhost:6397/0", queue_configs=[QueueConfig(name="samples")]))
+app = Litestar(plugins=[saq], route_handlers=[SampleController])

examples/basic.py

@@ -0,0 +1,26 @@
+import asyncio
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+async def system_upkeep(_: dict) -> None:
+    logger.info("Performing system upkeep operations.")
+    logger.info("Simulating a long running operation.  Sleeping for 60 seconds.")
+    await asyncio.sleep(60)
+    logger.info("Simulating an even longer running operation.  Sleeping for 120 seconds.")
+    await asyncio.sleep(120)
+    logger.info("Long running process complete.")
+    logger.info("Performing system upkeep operations.")
+
+
+async def background_worker_task(_: dict) -> None:
+    logger.info("Performing background worker task.")
+    await asyncio.sleep(20)
+    logger.info("Performing system upkeep operations.")
+
+
+async def system_task(_: dict) -> None:
+    logger.info("Performing simple system task")
+    await asyncio.sleep(2)
+    logger.info("System task complete.")

examples/tasks.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
 from . import controllers, info
-from .base import BackgroundTaskError, CronJob, Job, Queue, Worker, WorkerFunction
-from .commands import create_worker_instance
+from .base import CronJob, Job, Queue, Worker
+from .config import QueueConfig, SAQConfig
+from .controllers import SAQController
+from .plugin import SAQPlugin
 
 __all__ = [
+    "SAQPlugin",
+    "SAQConfig",
+    "SAQController",
+    "QueueConfig",
     "Queue",
     "CronJob",
     "Job",
     "Worker",
-    "WorkerFunction",
-    "create_worker_instance",
-    "BackgroundTaskError",
     "info",
     "controllers",
 ]

litestar_saq/__init__.py

@@ -1,40 +1,39 @@
 from __future__ import annotations
 
 import asyncio
-from collections import abc
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-import saq
+from saq import Job as SaqJob
+from saq import Worker as SaqWorker
+from saq.job import CronJob as SaqCronJob
+from saq.queue import Queue as SaqQueue
 
 if TYPE_CHECKING:
+    from collections.abc import Collection
     from signal import Signals
 
-
-WorkerFunction = abc.Callable[..., abc.Awaitable[Any]]
+    from redis.asyncio.client import Redis
+    from saq.types import DumpType, Function, LoadType, PartialTimersDict, ReceivesContext
 
 
 @dataclass
-class Job(saq.Job):
+class Job(SaqJob):
     """Job Details"""
 
     job_name: str | None = None
     job_description: str | None = None
 
 
 @dataclass
-class CronJob(saq.CronJob):
+class CronJob(SaqCronJob):
     """Cron Job Details"""
 
     job_name: str | None = None
     job_description: str | None = None
 
 
-class BackgroundTaskError(Exception):
-    """Base class for `Task` related exceptions."""
-
-
-class Queue(saq.Queue):
+class Queue(SaqQueue):
     """[SAQ Queue](https://github.com/tobymao/saq/blob/master/saq/queue.py).
 
     Configures `msgspec` for msgpack serialization/deserialization if not otherwise configured.
@@ -47,22 +46,72 @@ class Queue(saq.Queue):
         Passed through to `saq.Queue.__init__()`
     """
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        redis: Redis[bytes],
+        name: str = "default",
+        dump: DumpType | None = None,
+        load: LoadType | None = None,
+        max_concurrent_ops: int = 20,
+        queue_namespace: str | None = None,
+    ) -> None:
+        self._namespace = queue_namespace if queue_namespace is not None else "saq"
+        super().__init__(redis, name, dump, load, max_concurrent_ops)
+
+    def temp(self, *args: Any, **kwargs: Any) -> None:
         """Initialize a new queue."""
-        """
-        kwargs.setdefault("dump", serialization.to_json)
-        kwargs.setdefault("load", serialization.from_json)
-        kwargs.setdefault("name", "background-tasks")
-        """
+        self._namespace = kwargs.pop("queue_namespace", "saq")
         super().__init__(*args, **kwargs)
 
+    def namespace(self, key: str) -> str:
+        """Make the namespace unique per app."""
+        return f"{self._namespace}:{self.name}:{key}"
+
+    def job_id(self, job_key: str) -> str:
+        """Job ID.
 
-class Worker(saq.Worker):
+        Args:
+            job_key (str): Sets the job ID for the given key
+
+        Returns:
+            str: Job ID for the specified key
+        """
+        return f"{self._namespace}:{self.name}:job:{job_key}"
+
+
+class Worker(SaqWorker):
     """Worker."""
 
     # same issue: https://github.com/samuelcolvin/arq/issues/182
     SIGNALS: list[Signals] = []
 
+    def __init__(
+        self,
+        queue: Queue | SaqQueue,
+        functions: Collection[Function | tuple[str, Function]],
+        *,
+        concurrency: int = 10,
+        cron_jobs: Collection[CronJob] | None = None,
+        startup: ReceivesContext | None = None,
+        shutdown: ReceivesContext | None = None,
+        before_process: ReceivesContext | None = None,
+        after_process: ReceivesContext | None = None,
+        timers: PartialTimersDict | None = None,
+        dequeue_timeout: float = 0,
+    ) -> None:
+        super().__init__(
+            queue,
+            functions,
+            concurrency=concurrency,
+            cron_jobs=cron_jobs,
+            startup=startup,
+            shutdown=shutdown,
+            before_process=before_process,
+            after_process=after_process,
+            timers=timers,
+            dequeue_timeout=dequeue_timeout,
+        )
+
     async def on_app_startup(self) -> None:
         """Attach the worker to the running event loop."""
         loop = asyncio.get_running_loop()

litestar_saq/base.py

@@ -1,24 +1,85 @@
 from __future__ import annotations
 
+import asyncio
+import multiprocessing
+from contextlib import suppress
 from typing import TYPE_CHECKING
 
-from click import group, option
+from click import IntRange, group, option
 from litestar.cli._utils import LitestarGroup, console
 
+from litestar_saq.exceptions import ImproperConfigurationError
+from litestar_saq.plugin import SAQPlugin
+
 if TYPE_CHECKING:
     from litestar import Litestar
 
+    from litestar_saq.base import Worker
+
 
-@group(cls=LitestarGroup, name="database")
+@group(cls=LitestarGroup, name="tasks")
 def background_worker_group() -> None:
     """Manage background task workers."""
 
 
 @background_worker_group.command(
-    name="show-current-revision",
-    help="Shows the current revision for the database.",
+    name="run-worker",
+    help="Run background worker processes.",
+)
+@option(
+    "--workers",
+    help="The number of worker processes to start.",
+    type=IntRange(min=1),
+    default=1,
+    required=False,
+    show_default=True,
 )
-@option("--verbose", type=bool, help="Enable verbose output.", default=False, is_flag=True)
-def show_database_revision(app: Litestar, verbose: bool) -> None:  # noqa: ARG001
-    """Show current database revision."""
+@option("-v", "--verbose", help="Enable verbose logging.", is_flag=True, default=None, type=bool, required=False)
+@option("-d", "--debug", help="Enable debugging.", is_flag=True, default=None, type=bool, required=False)
+def run_worker(
+    app: Litestar,
+    workers: int,
+    verbose: bool | None,  # noqa: ARG001
+    debug: bool | None,  # noqa: ARG001
+) -> None:
+    """Run the API server."""
     console.rule("[yellow]Starting SAQ Workers[/]", align="left")
+
+    plugin = get_saq_plugin(app)
+    if workers > 1:
+        for _ in range(workers - 1):
+            p = multiprocessing.Process(target=run_worker_process, args=(plugin.get_workers(),))
+            p.start()
+
+    with suppress(KeyboardInterrupt):
+        run_worker_process(workers=plugin.get_workers())
+
+
+def get_saq_plugin(app: Litestar) -> SAQPlugin:
+    """Retrieve a SAQ plugin from the Litestar application's plugins.
+
+    This function attempts to find a SAQ plugin instance.
+    If plugin is not found, it raises an ImproperlyConfiguredException.
+    """
+
+    with suppress(KeyError):
+        return app.plugins.get(SAQPlugin)
+    msg = "Failed to initialize SAQ. The required plugin (SAQPlugin) is missing."
+    raise ImproperConfigurationError(
+        msg,
+    )
+
+
+def run_worker_process(workers: list[Worker]) -> None:
+    """Run a worker."""
+    loop = asyncio.get_event_loop()
+
+    try:
+        for i, worker_instance in enumerate(workers):
+            if i < len(workers) - 1:
+                loop.create_task(worker_instance.start())
+            else:
+                loop.run_until_complete(worker_instance.start())
+    except KeyboardInterrupt:
+        for worker in workers:
+            loop.run_until_complete(worker.stop())