oban 0.5.0__py3-none-any.whl

oban/cli.py

@@ -0,0 +1,436 @@
+from __future__ import annotations
+
+import asyncio
+import importlib
+import logging
+import os
+import signal
+import socket
+import subprocess
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+import click
+import orjson
+from psycopg.types.json import set_json_dumps, set_json_loads
+from psycopg_pool import AsyncConnectionPool
+
+from oban import __version__
+from oban._config import Config
+from oban.schema import (
+    install as install_schema,
+    uninstall as uninstall_schema,
+)
+from oban.telemetry import logger as telemetry_logger
+
+try:
+    from uvloop import run as asyncio_run
+except ImportError:
+    from asyncio import run as asyncio_run
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+logger = logging.getLogger("oban.cli")
+
+
+def _file_to_module(file_path: str) -> str | None:
+    root = Path(os.getcwd())
+    path = Path(file_path)
+
+    try:
+        rel_path = path.relative_to(root)
+    except ValueError:
+        return None
+
+    parts = list(rel_path.parts)
+
+    if parts[-1].endswith(".py"):
+        parts[-1] = parts[-1][:-3]
+
+    if parts[-1] == "__init__":
+        parts.pop()
+
+    return ".".join(parts)
+
+
+def _import_cron_modules(module_paths: list[str]) -> int:
+    def safe_import(path: str) -> bool:
+        try:
+            importlib.import_module(path)
+            return True
+        except Exception as error:
+            logger.error(f"Failed to import cron module at {path}: {error}")
+
+            return False
+
+    return sum([safe_import(path) for path in module_paths])
+
+
+def _import_cron_paths(paths: list[str]) -> list[str]:
+    root = Path(os.getcwd())
+    grep = ["grep", "-rl", "--include=*.py", r"@worker.*cron\|@job.*cron"]
+
+    found = []
+    for pattern in [str(root / path) for path in paths]:
+        result = subprocess.run(
+            [*grep, pattern],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+
+        if result.returncode == 0:
+            files = [
+                line.strip()
+                for line in result.stdout.strip().split("\n")
+                if line.strip()
+            ]
+            found.extend(files)
+
+    files = set(str(Path(file).resolve()) for file in found)
+
+    return [mod for file in files if (mod := _file_to_module(file))]
+
+
+def _split_csv(value: str | None) -> list[str] | None:
+    return [item.strip() for item in value.split(",")] if value else None
+
+
+def _find_and_load_cron_modules(
+    cron_modules: list[str] | None = None, cron_paths: list[str] | None = None
+) -> None:
+    if cron_modules:
+        logger.info(f"Importing {len(cron_modules)} cron modules...")
+
+    elif cron_paths:
+        logger.info(f"Discovering cron workers in {', '.join(cron_paths)}...")
+
+        cron_modules = _import_cron_paths(cron_paths)
+    else:
+        logger.info("Auto-discovering cron workers in current directory...")
+
+        cron_modules = _import_cron_paths([os.getcwd()])
+
+    import_count = _import_cron_modules(cron_modules)
+
+    logger.info(
+        f"Successfully imported {import_count}/{len(cron_modules)} cron modules"
+    )
+
+
+def print_banner(version: str) -> None:
+    banner = f"""
+  
+   ██████╗  ██████╗    ████╗   ███╗   ██╗
+  ██╔═══██╗ ██╔══██╗  ██╔═██╗  ████╗  ██║
+  ██║   ██║ ██████╔╝ ████████╗ ██╔██╗ ██║
+  ██║   ██║ ██╔══██╗ ██╔═══██║ ██║╚██╗██║
+   ██████╔╝ ██████╔╝ ██║   ██║ ██║ ╚████║
+   ╚═════╝  ╚═════╝  ╚═╝   ╚═╝ ╚═╝  ╚═══╝
+  
+  Job orchestration framework for Python, backed by PostgreSQL
+
+  v{version} | https://oban.pro
+"""
+    print(banner)
+
+
+@asynccontextmanager
+async def schema_pool(dsn: str) -> AsyncIterator[AsyncConnectionPool]:
+    if not dsn:
+        raise click.UsageError("--dsn is required (or set OBAN_DSN)")
+
+    conf = Config(dsn=dsn, pool_min_size=1, pool_max_size=1)
+    pool = await conf.create_pool()
+
+    try:
+        yield pool
+    finally:
+        await pool.close()
+
+
+async def _start_pool(conf: Config) -> AsyncConnectionPool:
+    try:
+        pool = await conf.create_pool()
+        logger.info(
+            f"Connected to database (pool: {conf.pool_min_size}-{conf.pool_max_size})"
+        )
+        return pool
+    except Exception as error:
+        logger.error(f"Failed to connect to database: {error!r}")
+        sys.exit(1)
+
+
+def handle_signals() -> asyncio.Event:
+    shutdown_event = asyncio.Event()
+    sigint_count = 0
+
+    def signal_handler(signum: int) -> None:
+        nonlocal sigint_count
+
+        shutdown_event.set()
+
+        if signum == signal.SIGTERM:
+            logger.info("Received SIGTERM, initiating graceful shutdown...")
+        elif signum == signal.SIGINT:
+            sigint_count += 1
+
+            if sigint_count == 1:
+                logger.info("Received SIGINT, initiating graceful shutdown...")
+                logger.info("Send another SIGINT to force exit")
+            else:
+                logger.warning("Forcing exit...")
+                sys.exit(1)
+
+    loop = asyncio.get_running_loop()
+    loop.add_signal_handler(signal.SIGTERM, lambda: signal_handler(signal.SIGTERM))
+    loop.add_signal_handler(signal.SIGINT, lambda: signal_handler(signal.SIGINT))
+
+    return shutdown_event
+
+
+@click.group(
+    context_settings={
+        "help_option_names": ["-h", "--help"],
+    }
+)
+@click.version_option(package_name="oban")
+def main() -> None:
+    """Oban - Job orchestration framework for Python, backed by PostgreSQL."""
+    pass
+
+
+@main.command()
+def version() -> None:
+    click.echo(f"oban {__version__}")
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True, dir_okay=False, path_type=str),
+    help="Path to TOML configuration file (default: searches for oban.toml)",
+)
+@click.option(
+    "--dsn",
+    envvar="OBAN_DSN",
+    help="PostgreSQL connection string",
+)
+@click.option(
+    "--prefix",
+    envvar="OBAN_PREFIX",
+    help="PostgreSQL schema name (default: public)",
+)
+def install(config: str | None, dsn: str | None, prefix: str | None) -> None:
+    """Install the Oban database schema."""
+
+    async def run() -> None:
+        conf = _load_conf(config, {"dsn": dsn, "prefix": prefix})
+        schema_prefix = conf.prefix or "public"
+
+        logger.info(f"Installing Oban schema in '{schema_prefix}'...")
+
+        try:
+            async with schema_pool(conf.dsn) as pool:
+                await install_schema(pool, prefix=schema_prefix)
+            logger.info("Schema installed successfully")
+        except Exception as error:
+            logger.error(f"Failed to install schema: {error!r}", exc_info=True)
+            sys.exit(1)
+
+    asyncio_run(run())
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True, dir_okay=False, path_type=str),
+    help="Path to TOML configuration file (default: searches for oban.toml)",
+)
+@click.option(
+    "--dsn",
+    envvar="OBAN_DSN",
+    help="PostgreSQL connection string",
+)
+@click.option(
+    "--prefix",
+    envvar="OBAN_PREFIX",
+    help="PostgreSQL schema name (default: public)",
+)
+def uninstall(config: str | None, dsn: str | None, prefix: str | None) -> None:
+    """Uninstall the Oban database schema."""
+
+    async def run() -> None:
+        conf = _load_conf(config, {"dsn": dsn, "prefix": prefix})
+        schema_prefix = conf.prefix or "public"
+
+        logger.info(f"Uninstalling Oban schema from '{schema_prefix}' schema...")
+
+        try:
+            async with schema_pool(conf.dsn) as pool:
+                await uninstall_schema(pool, prefix=schema_prefix)
+            logger.info("Schema uninstalled successfully")
+        except Exception as e:
+            logger.error(f"Failed to uninstall schema: {e}", exc_info=True)
+            sys.exit(1)
+
+    asyncio_run(run())
+
+
+@main.command()
+@click.option(
+    "--config",
+    type=click.Path(exists=True, dir_okay=False, path_type=str),
+    help="Path to TOML configuration file (default: searches for oban.toml)",
+)
+@click.option(
+    "--dsn",
+    envvar="OBAN_DSN",
+    help="PostgreSQL connection string",
+)
+@click.option(
+    "--queues",
+    envvar="OBAN_QUEUES",
+    help="Comma-separated queue:limit pairs (e.g., 'default:10,mailers:5')",
+)
+@click.option(
+    "--prefix",
+    envvar="OBAN_PREFIX",
+    help="PostgreSQL schema name (default: public)",
+)
+@click.option(
+    "--node",
+    envvar="OBAN_NODE",
+    help="Node identifier (default: hostname)",
+)
+@click.option(
+    "--pool-min-size",
+    envvar="OBAN_POOL_MIN_SIZE",
+    type=int,
+    help="Minimum connection pool size (default: 1)",
+)
+@click.option(
+    "--pool-max-size",
+    envvar="OBAN_POOL_MAX_SIZE",
+    type=int,
+    help="Maximum connection pool size (default: 10)",
+)
+@click.option(
+    "--log-level",
+    type=click.Choice(
+        ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False
+    ),
+    default="INFO",
+    help="Logging level (default: INFO)",
+)
+@click.option(
+    "--cron-modules",
+    envvar="OBAN_CRON_MODULES",
+    help="Comma-separated list of module paths with cron workers (e.g., 'myapp.workers,myapp.jobs')",
+)
+@click.option(
+    "--cron-paths",
+    envvar="OBAN_CRON_PATHS",
+    help="Comma-separated list of directories to search for cron workers (e.g., 'myapp/workers')",
+)
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    help="Validate configuration and load cron modules without starting Oban",
+)
+def start(
+    log_level: str,
+    config: str | None,
+    cron_modules: str | None,
+    cron_paths: str | None,
+    dry_run: bool,
+    **params: Any,
+) -> None:
+    """Start the Oban worker process.
+
+    This command starts an Oban instance that processes jobs from the configured queues.
+    The process will run until terminated by a signal.
+
+    Signal handling:
+    - SIGTERM: Graceful shutdown (finish running jobs, then exit)
+    - SIGINT (Ctrl+C): Graceful shutdown on first signal, force exit on second
+
+    Examples:
+
+        # Start with queues
+        oban start --dsn postgresql://localhost/mydb --queues default:10,mailers:5
+
+        # Use environment variables
+        export OBAN_DSN=postgresql://localhost/mydb
+        export OBAN_QUEUES=default:10,mailers:5
+        oban start
+    """
+    logging.getLogger().setLevel(getattr(logging, log_level.upper()))
+
+    set_json_dumps(orjson.dumps)
+    set_json_loads(orjson.loads)
+
+    conf = _load_conf(config, params)
+    node = conf.node or socket.gethostname()
+
+    async def run() -> None:
+        print_banner(__version__)
+
+        logger.info(f"Starting Oban v{__version__} on node {node}...")
+
+        _find_and_load_cron_modules(
+            cron_modules=_split_csv(cron_modules),
+            cron_paths=_split_csv(cron_paths),
+        )
+
+        if dry_run:
+            logger.info("Dry run complete-configuration is valid!")
+            sys.exit(0)
+
+        pool = await _start_pool(conf)
+        oban = await conf.create_oban(pool)
+
+        telemetry_logger.attach()
+        shutdown_event = handle_signals()
+
+        try:
+            async with oban:
+                logger.info("Oban started, press Ctrl+C to stop")
+
+                await shutdown_event.wait()
+
+                logger.info("Shutting down gracefully...")
+        except Exception as error:
+            logger.error(f"Error during operation: {error!r}", exc_info=True)
+            sys.exit(1)
+        finally:
+            telemetry_logger.detach()
+            await pool.close()
+            logger.info("Shutdown complete")
+
+    asyncio_run(run())
+
+
+def _load_conf(conf_path: str | None, params: Any) -> Config:
+    if conf_path and not Path(conf_path).exists():
+        raise click.UsageError(f"--config file '{conf_path}' doesn't exist")
+
+    if queues := params.pop("queues", None):
+        params["queues"] = Config._parse_queues(queues)
+
+    conf = Config.load(conf_path, **params)
+
+    if not conf.dsn:
+        raise click.UsageError("--dsn, OBAN_DSN, or dsn in oban.toml required")
+
+    return conf
+
+
+if __name__ == "__main__":
+    main()

oban/decorators.py

@@ -0,0 +1,218 @@
+"""
+Decorators for creating Oban workers and jobs.
+
+This module provides two decorators for making your code enqueueable:
+
+- `@worker` For classes with a `process` method
+- `@job` For wrapping functions as jobs
+"""
+
+import inspect
+from functools import wraps
+from typing import Any, Callable
+
+from ._extensions import use_ext
+from ._worker import register_worker, worker_name
+from ._scheduler import register_scheduled
+from .job import Job, Result
+
+JOB_FIELDS = set(Job.__dataclass_fields__.keys()) - {"extra"} | {"schedule_in"}
+
+
+def worker(*, oban: str = "oban", cron: str | dict | None = None, **overrides):
+    """Decorate a class to make it a viable worker.
+
+    The decorator adds worker functionality to a class, including job creation
+    and enqueueing methods. The decorated class must implement a `process` method.
+
+    For simpler function-based jobs, consider using @job instead.
+
+    Args:
+        oban: Name of the Oban instance to use (default: "oban")
+        cron: Optional cron configuration for periodic execution. Can be:
+              - A string expression (e.g., "0 0 \\* \\* \\*" or "@daily")
+              - A dict with "expr" and optional "timezone" keys (timezone as string)
+        **overrides: Configuration options for the worker (queue, priority, etc.)
+
+    Returns:
+        A decorator function that can be applied to worker classes
+
+    Example:
+        >>> from oban import Oban, worker
+        >>>
+        >>> # Create an Oban instance with a specific name
+        >>> oban_instance = Oban(name="oban", queues={"default": 10, "mailers": 5})
+        >>>
+        >>> @worker(queue="mailers", priority=1)
+        ... class EmailWorker:
+        ...     async def process(self, job):
+        ...         print(f"Sending email: {job.args}")
+        ...         return None
+        >>>
+        >>> # Create a job without enqueueing
+        >>> job = EmailWorker.new({"to": "user@example.com", "subject": "Hello"})
+        >>> print(job.queue)  # "mailers"
+        >>> print(job.priority)  # 1
+        >>>
+        >>> # Create and enqueue a job
+        >>> job = EmailWorker.enqueue(
+        ...     {"to": "admin@example.com", "subject": "Alert"},
+        ...     priority=5  # Override default priority
+        ... )
+        >>> print(job.priority)  # 5
+        >>>
+        >>> # Schedule a job to run in 5 minutes using a timedelta
+        >>> from datetime import timedelta
+        >>> job = EmailWorker.enqueue(..., schedule_in=timedelta(minutes=5))
+        >>>
+        >>> # Schedule a job to run in 60 seconds
+        >>> job = EmailWorker.enqueue(...,schedule_in=60)
+        >>>
+        >>> # Periodic worker that runs daily at midnight
+        >>> @worker(queue="cleanup", cron="@daily")
+        ... class DailyCleanup:
+        ...     async def process(self, job):
+        ...         print("Running daily cleanup")
+        ...         return None
+        >>>
+        >>> # Periodic worker with timezone
+        >>> @worker(queue="reports", cron={"expr": "0 9 \\* \\* MON-FRI", "timezone": "America/New_York"})
+        ... class BusinessHoursReport:
+        ...     async def process(self, job):
+        ...         print("Running during NY business hours")
+        ...         return None
+        >>>
+        >>> # Workers can also be created without args
+        >>> job = DailyCleanup.new()  # args defaults to {}
+        >>>
+        >>> # Custom backoff for retries
+        >>> @worker(queue="default")
+        ... class CustomBackoffWorker:
+        ...     async def process(self, job):
+        ...         return None
+        ...
+        ...     def backoff(self, job):
+        ...         # Simple linear backoff at 2x the attempt number
+        ...         return 2 * job.attempt
+
+    Note:
+        The worker class must implement a ``process(self, job: Job) -> Result[Any]`` method.
+        If not implemented, a NotImplementedError will be raised when called.
+
+        Optionally implement a ``backoff(self, job: Job) -> int`` method to customize
+        retry delays. If not provided, uses Oban's default jittery clamped backoff.
+    """
+
+    def decorate(cls: type) -> type:
+        if not hasattr(cls, "process"):
+
+            async def process(self, job: Job) -> Result[Any]:
+                raise NotImplementedError("Worker must implement process method")
+
+            setattr(cls, "process", process)
+
+        @classmethod
+        def new(cls, args: dict[str, Any] | None = None, /, **params) -> Job:
+            merged = {**cls._opts, **params}
+            extras = {
+                key: merged.pop(key) for key in list(merged) if key not in JOB_FIELDS
+            }
+
+            if extras:
+                merged["extra"] = extras
+
+            return Job.new(worker=worker_name(cls), args=args or {}, **merged)
+
+        @classmethod
+        async def enqueue(
+            cls, args: dict[str, Any] | None = None, /, **overrides
+        ) -> Job:
+            from .oban import get_instance
+
+            job = cls.new(args, **overrides)
+
+            return await get_instance(cls._oban_name).enqueue(job)
+
+        setattr(cls, "_opts", overrides)
+        setattr(cls, "_oban_name", oban)
+        setattr(cls, "new", new)
+        setattr(cls, "enqueue", enqueue)
+
+        register_worker(cls)
+
+        use_ext("worker.after_register", lambda _cls: None, cls)
+
+        if cron:
+            register_scheduled(cron, cls)
+
+        return cls
+
+    return decorate
+
+
+def job(*, oban: str = "oban", cron: str | dict | None = None, **overrides):
+    """Decorate a function to make it an Oban job.
+
+    The decorated function's signature is preserved for new() and enqueue().
+
+    Use @job for simple function-based tasks where you don't need access to
+    job metadata such as the attempt, past errors.
+
+    Args:
+        oban: Name of the Oban instance to use (default: "oban")
+        cron: Optional cron configuration for periodic execution. Can be:
+              - A string expression (e.g., "0 0 \\* \\* \\*" or "@daily")
+              - A dict with "expr" and optional "timezone" keys (timezone as string)
+        **overrides: Configuration options (queue, priority, etc.)
+
+    Example:
+        >>> from oban import job
+        >>>
+        >>> @job(queue="mailers", priority=1)
+        ... def send_email(to: str, subject: str, body: str):
+        ...     print(f"Sending to {to}: {subject}")
+        >>>
+        >>> send_email.enqueue("user@example.com", "Hello", "World")
+        >>>
+        >>> # Periodic job that runs weekly at midnight
+        >>> @job(queue="reports", cron="@weekly")
+        ... def generate_weekly_report():
+        ...     print("Generating weekly report")
+        ...     return {"status": "complete"}
+    """
+
+    def decorate(func: Callable[..., Any]) -> type:
+        sig = inspect.signature(func)
+
+        class FunctionWorker:
+            async def process(self, job: Job):
+                return func(**job.args)
+
+        FunctionWorker.__name__ = func.__name__  # type: ignore[attr-defined]
+        FunctionWorker.__module__ = func.__module__  # type: ignore[attr-defined]
+        FunctionWorker.__qualname__ = func.__qualname__  # type: ignore[attr-defined]
+        FunctionWorker.__doc__ = func.__doc__
+
+        worker_cls = worker(oban=oban, cron=cron, **overrides)(FunctionWorker)
+
+        original_new = worker_cls.new
+        original_enq = worker_cls.enqueue
+
+        @wraps(func)
+        def new_with_sig(*args, **kwargs):
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            return original_new(dict(bound.arguments))
+
+        @wraps(func)
+        async def enq_with_sig(*args, **kwargs):
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            return await original_enq(dict(bound.arguments))
+
+        worker_cls.new = staticmethod(new_with_sig)
+        worker_cls.enqueue = staticmethod(enq_with_sig)
+
+        return worker_cls
+
+    return decorate