SAQ Plugin for Litestar
from __future__ import annotations

from typing import TYPE_CHECKING

from litestar import Controller, Litestar, get

from examples import tasks
from litestar_saq import QueueConfig, SAQConfig, SAQPlugin
from litestar_saq.base import CronJob

if TYPE_CHECKING:
    from saq.types import QueueInfo

    from litestar_saq.config import TaskQueues


class SampleController(Controller):
    @get(path="/samples")
    async def samples_queue_info(self, task_queues: TaskQueues) -> QueueInfo:
        """Check database available and returns app config info."""
        queue = task_queues.get("samples")
        return await queue.info()


saq = SAQPlugin(
    config=SAQConfig(
        redis_url="redis://localhost:6397/0",
        web_enabled=True,
        queue_configs=[
            QueueConfig(
                name="samples",
                tasks=[tasks.background_worker_task, tasks.system_task, tasks.system_upkeep],
                scheduled_tasks=[CronJob(function=tasks.system_upkeep, cron="* * * * *", timeout=600, ttl=2000)],
            ),
        ],
    ),
)
app = Litestar(plugins=[saq], route_handlers=[SampleController])

examples/basic.py
cofin/litestar-saq
from __future__ import annotations

import asyncio
from logging import getLogger
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from saq.types import Context

logger = getLogger(__name__)


async def system_upkeep(_: Context) -> 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(_: Context) -> None:
    logger.info("Performing background worker task.")
    await asyncio.sleep(20)
    logger.info("Performing system upkeep operations.")


async def system_task(_: Context) -> None:
    logger.info("Performing simple system task")
    await asyncio.sleep(2)
    logger.info("System task complete.")

examples/tasks.py
cofin/litestar-saq
from __future__ import annotations

from .base import CronJob, Job, Queue, Worker
from .config import QueueConfig, SAQConfig
from .plugin import SAQPlugin

__all__ = (
    "SAQPlugin",
    "SAQConfig",
    "QueueConfig",
    "Queue",
    "CronJob",
    "Job",
    "Worker",
)

litestar_saq/__init__.py
cofin/litestar-saq
"""Metadata for the Project."""
from __future__ import annotations

import importlib.metadata

__all__ = ["__version__", "__project__"]

__version__ = importlib.metadata.version("litestar_saq")
"""Version of the project."""
__project__ = importlib.metadata.metadata("litestar_saq")["Name"]
"""Name of the project."""

litestar_saq/__metadata__.py
cofin/litestar-saq
from __future__ import annotations

import pkgutil
import sys
from functools import lru_cache
from importlib import import_module
from importlib.machinery import SourceFileLoader
from pathlib import Path
from typing import TYPE_CHECKING, Any

__all__ = ["import_string", "module_to_os_path"]


if TYPE_CHECKING:
    from types import ModuleType


@lru_cache
def module_to_os_path(dotted_path: str) -> Path:
    """Find Module to OS Path.

    Return path to the base directory of the project or the module
    specified by `dotted_path`.

    Ensures that pkgutil returns a valid source file loader.
    """
    src = pkgutil.get_loader(dotted_path)
    if not isinstance(src, SourceFileLoader):
        msg = f"Couldn't find the path for {dotted_path}"
        raise TypeError(msg)
    return Path(str(src.path).removesuffix("/__init__.py"))


def import_string(dotted_path: str) -> Any:
    """Dotted Path Import.

    Import a dotted module path and return the attribute/class designated by the
    last name in the path. Raise ImportError if the import failed.

    Args:
        dotted_path (str): The path of the module to import.

    Raises:
        ImportError: Could not import the module.

    Returns:
        object: The imported object.
    """
    try:
        module_path, class_name = dotted_path.rsplit(".", 1)
    except ValueError as e:
        msg = f"{dotted_path} doesn't look like a module path"
        raise ImportError(msg) from e

    try:
        return _cached_import(module_path, class_name)
    except AttributeError as e:
        msg = f"Module '{module_path}' does not define a '{class_name}' attribute/class"
        raise ImportError(msg) from e


def _is_loaded(module: ModuleType | None) -> bool:
    spec = getattr(module, "__spec__", None)
    initializing = getattr(spec, "_initializing", False)
    return bool(module and spec and not initializing)


def _cached_import(module_path: str, class_name: str) -> Any:
    """Import and cache a class from a module.

    Args:
        module_path (str): dotted path to module.
        class_name (str): Class or function name.

    Returns:
        object: The imported class or function
    """
    # Check whether module is loaded and fully initialized.
    module = sys.modules.get(module_path)
    if not _is_loaded(module):
        module = import_module(module_path)
    return getattr(module, class_name)

litestar_saq/_util.py
cofin/litestar-saq
from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast

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

from litestar_saq._util import import_string

if TYPE_CHECKING:
    from collections.abc import Collection
    from signal import Signals

    from redis.asyncio.client import Redis
    from saq.types import DumpType, Function, LoadType, PartialTimersDict, ReceivesContext


@dataclass
class Job(SaqJob):
    """Job Details"""


@dataclass
class CronJob(SaqCronJob):
    """Cron Job Details"""

    function: Function | str  # type: ignore[assignment]
    meta: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.function = self._get_or_import_function(self.function)

    @staticmethod
    def _get_or_import_function(function_or_import_string: str | Function) -> Function:
        if isinstance(function_or_import_string, str):
            return cast("Function", import_string(function_or_import_string))
        return function_or_import_string


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.

    Parameters
    ----------
    *args : Any
        Passed through to `saq.Queue.__init__()`
    **kwargs : Any
        Passed through to `saq.Queue.__init__()`
    """

    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 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.

        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()
        _ = loop.create_task(self.start())

litestar_saq/base.py
cofin/litestar-saq
from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from click import Group
    from litestar import Litestar
    from litestar.logging.config import BaseLoggingConfig

    from litestar_saq.base import Worker
    from litestar_saq.plugin import SAQPlugin


def build_cli_app() -> Group:
    import asyncio
    import multiprocessing
    from typing import cast

    from click import IntRange, group, option
    from litestar.cli._utils import LitestarGroup, console

    @group(cls=LitestarGroup, name="workers")
    def background_worker_group() -> None:
        """Manage background task workers."""

    @background_worker_group.command(
        name="run",
        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("-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,
        debug: bool | None,
    ) -> None:
        """Run the API server."""
        console.rule("[yellow]Starting SAQ Workers[/]", align="left")
        if app.logging_config is not None:
            app.logging_config.configure()
        if debug is not None or verbose is not None:
            app.debug = True
        plugin = get_saq_plugin(app)
        show_saq_info(app, workers, plugin)
        if workers > 1:
            for _ in range(workers - 1):
                p = multiprocessing.Process(target=run_worker_process, args=(plugin.get_workers(), app.logging_config))
                p.start()

        try:
            run_worker_process(
                workers=plugin.get_workers(),
                logging_config=cast("BaseLoggingConfig", app.logging_config),
            )
        except KeyboardInterrupt:
            loop = asyncio.get_event_loop()
            for worker_instance in plugin.get_workers():
                loop.run_until_complete(worker_instance.stop())

    return background_worker_group


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.
    """

    from contextlib import suppress

    from litestar_saq.exceptions import ImproperConfigurationError
    from litestar_saq.plugin import SAQPlugin

    with suppress(KeyError):
        return app.plugins.get(SAQPlugin)
    msg = "Failed to initialize SAQ. The required plugin (SAQPlugin) is missing."
    raise ImproperConfigurationError(
        msg,
    )


def show_saq_info(app: Litestar, workers: int, plugin: SAQPlugin) -> None:  # pragma: no cover
    """Display basic information about the application and its configuration."""

    from litestar.cli._utils import _format_is_enabled, console
    from rich.table import Table
    from saq import __version__ as saq_version

    table = Table(show_header=False)
    table.add_column("title", style="cyan")
    table.add_column("value", style="bright_blue")

    table.add_row("SAQ version", saq_version)
    table.add_row("Debug mode", _format_is_enabled(app.debug))
    table.add_row("Number of Processes", str(workers))
    table.add_row("Queues", str(len(plugin._config.queue_configs)))  # noqa: SLF001

    console.print(table)


def run_worker_process(workers: list[Worker], logging_config: BaseLoggingConfig | None) -> None:
    """Run a worker."""
    import asyncio

    loop = asyncio.get_event_loop()
    if logging_config is not None:
        logging_config.configure()
    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())

litestar_saq/cli.py
cofin/litestar-saq
from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, TypeVar, cast

from litestar.exceptions import ImproperlyConfiguredException
from litestar.serialization import decode_json, encode_json
from redis.asyncio import ConnectionPool, Redis
from saq.queue import Queue as SaqQueue
from saq.types import DumpType, LoadType, PartialTimersDict, QueueInfo, QueueStats, ReceivesContext

from litestar_saq._util import import_string, module_to_os_path
from litestar_saq.base import CronJob, Job, Queue, Worker

if TYPE_CHECKING:
    from typing import Any

    from litestar import Litestar
    from litestar.datastructures.state import State

T = TypeVar("T")


def serializer(value: Any) -> str:
    """Serialize JSON field values.

    Args:
        value: Any json serializable value.

    Returns:
        JSON string.
    """
    return encode_json(value).decode("utf-8")


def _get_static_files() -> Path:
    return Path(module_to_os_path("saq.web") / "static")


TaskQueue = SaqQueue | Queue


@dataclass
class TaskQueues:
    queues: dict[str, TaskQueue] = field(default_factory=dict)

    def get(self, name: str) -> TaskQueue:
        queue = self.queues.get(name)
        if queue is not None:
            return queue
        msg = "Could not find the specified queue.  Please check your configuration."
        raise ImproperlyConfiguredException(msg)


@dataclass
class SAQConfig:
    """SAQ Configuration."""

    queue_configs: list[QueueConfig] = field(default_factory=lambda: [QueueConfig()])
    """Configuration for Queues"""
    redis: Redis | None = None
    """Pre-configured Redis instance to use."""
    redis_url: str | None = None
    """Redis URL to connect with."""
    namespace: str = "saq"
    """Namespace to use for Redis"""
    queue_instances: dict[str, Queue | SaqQueue] | None = None
    """Current configured queue instances.  When None, queues will be auto-created on startup"""
    queues_dependency_key: str = field(default="task_queues")
    """Key to use for storing dependency information in litestar."""
    worker_processes: int = 1
    """The number of worker processes to spawn.

    Default is set to 1.
    """
    web_enabled: bool = False
    """If true, the worker admin UI is launched on worker startup.."""
    json_deserializer: LoadType = decode_json
    """This is a Python callable that will
    convert a JSON string to a Python object. By default, this is set to Litestar's
    :attr:`decode_json() <.serialization.decode_json>` function."""
    json_serializer: DumpType = serializer
    """This is a Python callable that will render a given object as JSON.
    By default, Litestar's :attr:`encode_json() <.serialization.encode_json>` is used."""
    static_files: Path = field(default_factory=_get_static_files)
    """Location of the static files to serve for the SAQ UI"""
    web_path = "/saq"
    """Base path to serve the SAQ web UI"""

    def __post_init__(self) -> None:
        if self.redis is not None and self.redis_url is not None:
            msg = "Only one of 'redis' or 'redis_url' can be provided."
            raise ImproperlyConfiguredException(msg)

    @property
    def signature_namespace(self) -> dict[str, Any]:
        """Return the plugin's signature namespace.

        Returns:
            A string keyed dict of names to be added to the namespace for signature forward reference resolution.
        """
        return {
            "Queue": Queue,
            "Worker": Worker,
            "QueueInfo": QueueInfo,
            "Job": Job,
            "QueueStats": QueueStats,
            "TaskQueues": TaskQueues,
        }

    async def on_shutdown(self, app: Litestar) -> None:
        """Disposes of the SAQ Workers.

        Args:
            app: The ``Litestar`` instance.

        Returns:
            None
        """

    def provide_queues(self, state: State) -> dict[str, Queue]:
        """Create an engine instance.

        Args:
            state: The ``Litestar.state`` instance.

        Returns:
            An engine instance.
        """
        return cast("dict[str, Queue]", state.get(self.queues_dependency_key))

    def get_redis(self) -> Redis:
        """Get the configured Redis connection.

        Returns:
            Dictionary of queues.
        """
        if self.redis is not None:
            return self.redis
        pool = ConnectionPool.from_url(
            url=cast("str", self.redis_url),
            decode_responses=False,
            socket_connect_timeout=2,
            socket_keepalive=5,
            health_check_interval=5,
        )
        self.redis = Redis(connection_pool=pool)
        return self.redis

    def get_queues(self) -> TaskQueues:
        """Get the configured SAQ queues.

        Returns:
            Dictionary of queues.
        """
        if self.queue_instances is not None:
            return TaskQueues(queues=self.queue_instances)
        self.queue_instances = {}
        for queue_config in self.queue_configs:
            self.queue_instances[queue_config.name] = Queue(
                queue_namespace=self.namespace,
                redis=self.get_redis(),
                name=queue_config.name,
                dump=self.json_serializer,
                load=self.json_deserializer,
                max_concurrent_ops=queue_config.max_concurrent_ops,
            )
        return TaskQueues(queues=self.queue_instances)

    def create_app_state_items(self) -> dict[str, Any]:
        """Key/value pairs to be stored in application state."""
        return {
            self.queues_dependency_key: self.get_queues(),
        }

    def update_app_state(self, app: Litestar) -> None:
        """Set the app state with worker queues.

        Args:
            app: The ``Litestar`` instance.
        """
        app.state.update(self.create_app_state_items())


@dataclass
class QueueConfig:
    """SAQ Queue Configuration"""

    redis: Redis | None = None
    """Pre-configured Redis instance to use."""
    name: str = "default"
    """The name of the queue to create."""
    concurrency: int = 10
    """Number of jobs to process concurrently"""
    max_concurrent_ops: int = 15
    """Maximum concurrent operations. (default 20)
            This throttles calls to `enqueue`, `job`, and `abort` to prevent the Queue
            from consuming too many Redis connections."""
    tasks: list[ReceivesContext | str] = field(default_factory=list)
    """Allowed list of functions to execute in this queue"""
    scheduled_tasks: list[CronJob] = field(default_factory=list)
    """Scheduled cron jobs to execute in this queue."""
    startup: ReceivesContext | None = None
    """Async callable to call on startup"""
    shutdown: ReceivesContext | None = None
    """Async callable to call on shutdown"""
    before_process: ReceivesContext | None = None
    """Async callable to call before a job processes"""
    after_process: ReceivesContext | None = None
    """Async callable to call after a job processes"""
    timers: PartialTimersDict | None = None
    """Dict with various timer overrides in seconds
            schedule: how often we poll to schedule jobs
            stats: how often to update stats
            sweep: how often to clean up stuck jobs
            abort: how often to check if a job is aborted"""
    dequeue_timeout: float = 0
    """How long it will wait to dequeue"""
    separate_process: bool = True
    """Executes as a separate event loop when True.
            Set it False to execute within the Litestar application."""

    def __post_init__(self) -> None:
        self.tasks = [self._get_or_import_task(task) for task in self.tasks]

    @staticmethod
    def _get_or_import_task(task_or_import_string: str | ReceivesContext) -> ReceivesContext:
        if isinstance(task_or_import_string, str):
            return cast("ReceivesContext", import_string(task_or_import_string))
        return task_or_import_string

litestar_saq/config.py
cofin/litestar-saq
from __future__ import annotations

from functools import lru_cache
from typing import TYPE_CHECKING, cast

from litestar.exceptions import NotFoundException

if TYPE_CHECKING:
    from litestar import Controller
    from saq.job import Job as SaqJob
    from saq.types import QueueInfo

    from litestar_saq.base import Job
    from litestar_saq.config import TaskQueue, TaskQueues


async def job_info(queue: TaskQueue, job_id: str) -> Job | SaqJob:
    job = await queue._get_job_by_id(job_id)  # noqa: SLF001
    if not job:
        msg = f"Could not find job ID {job_id}"
        raise NotFoundException(msg)
    return cast("Job", job)


@lru_cache
def build_controller(url_base: str = "/saq") -> type[Controller]:  # noqa: C901
    from litestar import Controller, MediaType, get, post
    from litestar.exceptions import NotFoundException
    from litestar.status_codes import HTTP_202_ACCEPTED

    class SAQController(Controller):
        tags = ["SAQ"]

        @get(
            operation_id="WorkerQueueList",
            name="worker:queue-list",
            path=f"{url_base}/api/queues/",
            media_type=MediaType.JSON,
            cache=False,
            summary="Queue List",
            description="List configured worker queues.",
        )
        async def queue_list(self, task_queues: TaskQueues) -> dict[str, list[QueueInfo]]:
            """Get Worker queues."""
            return {"queues": [await queue.info() for queue in task_queues.queues.values()]}

        @get(
            operation_id="WorkerQueueDetail",
            name="worker:queue-detail",
            path=f"{url_base}/api/queues/{{queue_id:str}}",
            media_type=MediaType.JSON,
            cache=False,
            summary="Queue Detail",
            description="List queue details.",
        )
        async def queue_detail(self, task_queues: TaskQueues, queue_id: str) -> dict[str, QueueInfo]:
            """Get queue information."""
            queue = task_queues.get(queue_id)
            if not queue:
                msg = f"Could not find the {queue_id} queue"
                raise NotFoundException(msg)
            return {"queue": await queue.info()}

        @get(
            operation_id="WorkerJobDetail",
            name="worker:job-detail",
            path=f"{url_base}/api/queues/{{queue_id:str}}/jobs/{{job_id:str}}",
            media_type=MediaType.JSON,
            cache=False,
            summary="Job Details",
            description="List job details.",
        )
        async def job_detail(self, task_queues: TaskQueues, queue_id: str, job_id: str) -> dict:
            """Get job information."""
            queue = task_queues.get(queue_id)
            if not queue:
                msg = f"Could not find the {queue_id} queue"
                raise NotFoundException(msg)
            job = await job_info(queue, job_id)
            job_dict = job.to_dict()
            if "kwargs" in job_dict:
                job_dict["kwargs"] = repr(job_dict["kwargs"])
            if "result" in job_dict:
                job_dict["result"] = repr(job_dict["result"])
            return {"job": job_dict}

        @post(
            operation_id="WorkerJobRetry",
            name="worker:job-retry",
            path=f"{url_base}/api/queues/{{queue_id:str}}/jobs/{{job_id:str}}/retry",
            media_type=MediaType.JSON,
            cache=False,
            summary="Job Retry",
            description="Retry a failed job..",
            status_code=HTTP_202_ACCEPTED,
        )
        async def job_retry(self, task_queues: TaskQueues, queue_id: str, job_id: str) -> dict:
            """Retry job."""
            queue = task_queues.get(queue_id)
            if not queue:
                msg = f"Could not find the {queue_id} queue"
                raise NotFoundException(msg)
            job = await job_info(queue, job_id)
            await job.retry("retried from ui")
            return {}

        @post(
            operation_id="WorkerJobAbort",
            name="worker:job-abort",
            path=f"{url_base}/api/queues/{{queue_id:str}}/jobs/{{job_id:str}}/abort",
            media_type=MediaType.JSON,
            cache=False,
            summary="Job Abort",
            description="Abort active job.",
            status_code=HTTP_202_ACCEPTED,
        )
        async def job_abort(self, task_queues: TaskQueues, queue_id: str, job_id: str) -> dict:
            """Abort job."""
            queue = task_queues.get(queue_id)
            if not queue:
                msg = f"Could not find the {queue_id} queue"
                raise NotFoundException(msg)
            job = await job_info(queue, job_id)
            await job.abort("aborted from ui")
            return {}

        # static site
        @get(
            [
                url_base,
                f"{url_base}/queues/{{queue_id:str}}",
                f"{url_base}/queues/{{queue_id:str}}/jobs/{{job_id:str}}",
            ],
            operation_id="WorkerIndex",
            name="worker:index",
            media_type=MediaType.HTML,
            cache=False,
            include_in_schema=False,
        )
        async def index(self) -> str:
            """Serve site root."""
            return f"""
            <!DOCTYPE html>
            <html>
                <head>
                    <meta charset="utf-8">
                    <meta name="viewport" content="width=device-width, initial-scale=1">
                    <link rel="stylesheet" type="text/css" href="{url_base}/static/pico.min.css.gz">
                    <title>SAQ</title>
                </head>
                <body>
                    <div id="app"></div>
                    <script>const root_path = "{url_base}";</script>
                    <script src="{url_base}/static/snabbdom.js.gz"></script>
                    <script src="{url_base}/static/app.js"></script>
                </body>
            </html>""".strip()

    return SAQController

litestar_saq/controllers.py
cofin/litestar-saq
class LitestarSaqError(Exception):
    """Base exception type for the Litestar Saq."""


class ImproperConfigurationError(LitestarSaqError):
    """Improper Configuration error.

    This exception is raised only when a module depends on a dependency that has not been installed.
    """


class BackgroundTaskError(Exception):
    """Base class for `Task` related exceptions."""

litestar_saq/exceptions.py
cofin/litestar-saq
from __future__ import annotations

from typing import TYPE_CHECKING, TypeVar

from litestar.di import Provide
from litestar.exceptions import ImproperlyConfiguredException
from litestar.plugins import CLIPluginProtocol, InitPluginProtocol
from litestar.static_files import StaticFilesConfig

from litestar_saq.base import Worker
from litestar_saq.controllers import build_controller

if TYPE_CHECKING:
    from click import Group
    from litestar.config.app import AppConfig

    from litestar_saq.config import SAQConfig, TaskQueue, TaskQueues


T = TypeVar("T")


class SAQPlugin(InitPluginProtocol, CLIPluginProtocol):
    """SAQ plugin."""

    __slots__ = ("_config",)

    def __init__(self, config: SAQConfig) -> None:
        """Initialize ``SAQPlugin``.

        Args:
            config: configure and start SAQ.
        """
        self._config = config
        self._worker_instances: list[Worker] | None = None

    def on_cli_init(self, cli: Group) -> None:
        from litestar_saq.cli import build_cli_app

        cli.add_command(build_cli_app())
        return super().on_cli_init(cli)

    def on_app_init(self, app_config: AppConfig) -> AppConfig:
        """Configure application for use with SQLAlchemy.

        Args:
            app_config: The :class:`AppConfig <.config.app.AppConfig>` instance.
        """
        app_config.dependencies.update(
            {
                self._config.queues_dependency_key: Provide(
                    dependency=self._config.provide_queues,
                    sync_to_thread=False,
                ),
            },
        )
        if self._config.web_enabled:
            app_config.static_files_config.append(
                StaticFilesConfig(
                    directories=[self._config.static_files],
                    path=f"{self._config.web_path}/static",
                    name="saq",
                    html_mode=False,
                    opt={"exclude_from_auth": True},
                ),
            )
            app_config.route_handlers.append(build_controller(self._config.web_path))
        app_config.on_startup.append(self._config.update_app_state)
        app_config.on_shutdown.append(self._config.on_shutdown)
        app_config.signature_namespace.update(self._config.signature_namespace)
        return app_config

    def get_workers(self) -> list[Worker]:
        """Return workers"""
        if self._worker_instances is not None:
            return self._worker_instances
        queues = self._config.get_queues()
        self._worker_instances = []
        self._worker_instances.extend(
            Worker(
                queue=self._get_queue(queue_config.name, queues.queues),
                functions=queue_config.tasks,
                cron_jobs=queue_config.scheduled_tasks,
                concurrency=queue_config.concurrency,
                startup=queue_config.startup,
                shutdown=queue_config.shutdown,
                before_process=queue_config.before_process,
                after_process=queue_config.after_process,
                timers=queue_config.timers,
                dequeue_timeout=queue_config.dequeue_timeout,
            )
            for queue_config in self._config.queue_configs
        )
        return self._worker_instances

    def get_queues(self) -> TaskQueues:
        return self._config.get_queues()

    @staticmethod
    def _get_queue(name: str, queues: dict[str, TaskQueue]) -> TaskQueue:
        queue = queues.get(name)
        if queue is not None:
            return queue
        msg = "Could not find the specified queue.  Please check your configuration."
        raise ImproperlyConfiguredException(msg)

litestar_saq/plugin.py
cofin/litestar-saq