python-pq 0.1.1__py3-none-any.whl

pq/__init__.py

@@ -0,0 +1,12 @@
+"""PQ - Postgres-backed task queue."""
+
+import pq.logging  # noqa: F401 - configures loguru on import
+
+from pq.client import PQ
+from pq.models import TaskStatus
+from pq.priority import Priority
+from pq.worker import TaskTimeoutError
+
+__version__ = "0.1.0"
+
+__all__ = ["PQ", "Priority", "TaskStatus", "TaskTimeoutError"]

pq/client.py

@@ -0,0 +1,352 @@
+"""PQ client - main interface for task queue."""
+
+from collections.abc import Callable, Set
+from contextlib import contextmanager
+from datetime import UTC, datetime, timedelta
+from typing import Any
+
+from croniter import croniter
+from croniter.croniter import CroniterBadCronError
+from sqlalchemy import create_engine, delete, func, select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker
+
+from pq.models import Base, Periodic, Task, TaskStatus
+from pq.priority import Priority
+from pq.registry import get_function_path
+from pq.serialization import serialize
+
+
+class PQ:
+    """Postgres-backed task queue client."""
+
+    def __init__(self, database_url: str) -> None:
+        """Initialize PQ with database connection.
+
+        Args:
+            database_url: PostgreSQL connection string.
+        """
+        self._engine: Engine = create_engine(database_url)
+        self._session_factory = sessionmaker(bind=self._engine)
+
+    @contextmanager
+    def session(self) -> Any:
+        """Get a database session context manager."""
+        session = self._session_factory()
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def create_tables(self) -> None:
+        """Create all tables (for testing)."""
+        Base.metadata.create_all(self._engine)
+
+    def drop_tables(self) -> None:
+        """Drop all tables (for testing)."""
+        Base.metadata.drop_all(self._engine)
+
+    def clear_all(self) -> None:
+        """Clear all tasks and periodic schedules."""
+        with self.session() as session:
+            session.execute(delete(Task))
+            session.execute(delete(Periodic))
+
+    def enqueue(
+        self,
+        task: Callable[..., Any],
+        *args: Any,
+        run_at: datetime | None = None,
+        priority: Priority = Priority.NORMAL,
+        **kwargs: Any,
+    ) -> int:
+        """Enqueue a one-off task.
+
+        Args:
+            task: Callable function to execute.
+            *args: Positional arguments to pass to the handler.
+            run_at: When to run the task. Defaults to now.
+            priority: Task priority. Higher = higher priority. Defaults to NORMAL.
+            **kwargs: Keyword arguments to pass to the handler.
+
+        Returns:
+            Task ID.
+
+        Raises:
+            ValueError: If task is a lambda, closure, or cannot be imported.
+        """
+        name = get_function_path(task)
+        payload = serialize(args, kwargs)
+
+        if run_at is None:
+            run_at = datetime.now(UTC)
+
+        task_obj = Task(name=name, payload=payload, run_at=run_at, priority=priority)
+
+        with self.session() as session:
+            session.add(task_obj)
+            session.flush()
+            return task_obj.id
+
+    def schedule(
+        self,
+        task: Callable[..., Any],
+        *args: Any,
+        run_every: timedelta | None = None,
+        cron: str | croniter | None = None,
+        priority: Priority = Priority.NORMAL,
+        **kwargs: Any,
+    ) -> int:
+        """Schedule a periodic task.
+
+        If a periodic task with this function already exists, it will be updated.
+        Either run_every or cron must be provided, but not both.
+
+        Args:
+            task: Callable function to execute.
+            *args: Positional arguments to pass to the handler.
+            run_every: Interval between executions (e.g., timedelta(hours=1)).
+            cron: Cron expression string (e.g., "0 9 * * 1") or croniter object.
+            priority: Task priority. Higher = higher priority. Defaults to NORMAL.
+            **kwargs: Keyword arguments to pass to the handler.
+
+        Returns:
+            Periodic task ID.
+
+        Raises:
+            ValueError: If neither run_every nor cron is provided, or if both are.
+            ValueError: If cron expression is invalid.
+            ValueError: If task is a lambda, closure, or cannot be imported.
+        """
+        if run_every is None and cron is None:
+            raise ValueError("Either run_every or cron must be provided")
+        if run_every is not None and cron is not None:
+            raise ValueError("Only one of run_every or cron can be provided")
+
+        # Validate and normalize cron expression
+        cron_expr: str | None = None
+        if cron is not None:
+            if isinstance(cron, croniter):
+                # Extract expression from croniter object
+                cron_expr = " ".join(str(f) for f in cron.expressions)
+            else:
+                # Validate string expression
+                try:
+                    croniter(cron)
+                except (KeyError, ValueError, CroniterBadCronError) as e:
+                    raise ValueError(f"Invalid cron expression '{cron}': {e}") from e
+                cron_expr = cron
+
+        name = get_function_path(task)
+        payload = serialize(args, kwargs)
+
+        # Calculate next_run based on cron or interval
+        now = datetime.now(UTC)
+        if cron_expr:
+            cron_iter = croniter(cron_expr, now)
+            next_run = cron_iter.get_next(datetime)
+        else:
+            next_run = now
+
+        with self.session() as session:
+            stmt = (
+                insert(Periodic)
+                .values(
+                    name=name,
+                    payload=payload,
+                    priority=priority,
+                    run_every=run_every,
+                    cron=cron_expr,
+                    next_run=next_run,
+                )
+                .on_conflict_do_update(
+                    index_elements=["name"],
+                    set_={
+                        "payload": payload,
+                        "priority": priority,
+                        "run_every": run_every,
+                        "cron": cron_expr,
+                        "next_run": next_run,
+                    },
+                )
+                .returning(Periodic.id)
+            )
+            result = session.execute(stmt)
+            return result.scalar_one()
+
+    def cancel(self, task_id: int) -> bool:
+        """Cancel a one-off task by ID.
+
+        Args:
+            task_id: Task ID.
+
+        Returns:
+            True if task was found and deleted, False otherwise.
+        """
+        with self.session() as session:
+            stmt = delete(Task).where(Task.id == task_id)
+            result = session.execute(stmt)
+            return result.rowcount > 0
+
+    def unschedule(self, task: Callable[..., Any]) -> bool:
+        """Remove a periodic task.
+
+        Args:
+            task: The scheduled function to remove.
+
+        Returns:
+            True if task was found and deleted, False otherwise.
+        """
+        name = get_function_path(task)
+        with self.session() as session:
+            stmt = delete(Periodic).where(Periodic.name == name)
+            result = session.execute(stmt)
+            return result.rowcount > 0
+
+    def pending_count(self) -> int:
+        """Count pending one-off tasks."""
+        with self.session() as session:
+            result = session.execute(
+                select(func.count())
+                .select_from(Task)
+                .where(Task.status == TaskStatus.PENDING)
+            )
+            return result.scalar_one()
+
+    def periodic_count(self) -> int:
+        """Count periodic task schedules."""
+        with self.session() as session:
+            result = session.execute(select(func.count()).select_from(Periodic))
+            return result.scalar_one()
+
+    def get_task(self, task_id: int) -> Task | None:
+        """Get a task by ID.
+
+        Args:
+            task_id: Task ID.
+
+        Returns:
+            Task object or None if not found.
+        """
+        with self.session() as session:
+            return session.get(Task, task_id)
+
+    def list_failed(self, limit: int = 100) -> list[Task]:
+        """List failed tasks.
+
+        Args:
+            limit: Maximum number of tasks to return.
+
+        Returns:
+            List of failed tasks, most recent first.
+        """
+        with self.session() as session:
+            stmt = (
+                select(Task)
+                .where(Task.status == TaskStatus.FAILED)
+                .order_by(Task.completed_at.desc())
+                .limit(limit)
+            )
+            return list(session.execute(stmt).scalars().all())
+
+    def list_completed(self, limit: int = 100) -> list[Task]:
+        """List completed tasks.
+
+        Args:
+            limit: Maximum number of tasks to return.
+
+        Returns:
+            List of completed tasks, most recent first.
+        """
+        with self.session() as session:
+            stmt = (
+                select(Task)
+                .where(Task.status == TaskStatus.COMPLETED)
+                .order_by(Task.completed_at.desc())
+                .limit(limit)
+            )
+            return list(session.execute(stmt).scalars().all())
+
+    def clear_completed(self, before: datetime | None = None) -> int:
+        """Clear completed tasks.
+
+        Args:
+            before: Only clear tasks completed before this time. If None, clears all.
+
+        Returns:
+            Number of tasks deleted.
+        """
+        with self.session() as session:
+            stmt = delete(Task).where(Task.status == TaskStatus.COMPLETED)
+            if before is not None:
+                stmt = stmt.where(Task.completed_at < before)
+            result = session.execute(stmt)
+            return result.rowcount
+
+    def clear_failed(self, before: datetime | None = None) -> int:
+        """Clear failed tasks.
+
+        Args:
+            before: Only clear tasks failed before this time. If None, clears all.
+
+        Returns:
+            Number of tasks deleted.
+        """
+        with self.session() as session:
+            stmt = delete(Task).where(Task.status == TaskStatus.FAILED)
+            if before is not None:
+                stmt = stmt.where(Task.completed_at < before)
+            result = session.execute(stmt)
+            return result.rowcount
+
+    def run_worker(
+        self,
+        *,
+        poll_interval: float = 1.0,
+        max_runtime: float = 30 * 60,
+        priorities: Set[Priority] | None = None,
+    ) -> None:
+        """Run the worker loop (blocking).
+
+        Each task executes in a forked child process for memory isolation.
+
+        Args:
+            poll_interval: Seconds to sleep between polls when idle.
+            max_runtime: Maximum execution time per task in seconds. Default: 30 min.
+            priorities: If set, only process tasks with these priority levels.
+                Use this to dedicate workers to specific priority tiers.
+        """
+        from pq.worker import run_worker
+
+        run_worker(
+            self,
+            poll_interval=poll_interval,
+            max_runtime=max_runtime,
+            priorities=priorities,
+        )
+
+    def run_worker_once(
+        self,
+        *,
+        max_runtime: float = 30 * 60,
+        priorities: Set[Priority] | None = None,
+    ) -> bool:
+        """Process a single task if available.
+
+        Each task executes in a forked child process for memory isolation.
+
+        Args:
+            max_runtime: Maximum execution time per task in seconds. Default: 30 min.
+            priorities: If set, only process tasks with these priority levels.
+
+        Returns:
+            True if a task was processed, False if queue was empty.
+        """
+        from pq.worker import run_worker_once
+
+        return run_worker_once(self, max_runtime=max_runtime, priorities=priorities)

pq/config.py

@@ -0,0 +1,11 @@
+"""PQ configuration using pydantic-settings."""
+
+from pydantic_settings import BaseSettings
+
+
+class PQSettings(BaseSettings):
+    """Configuration for PQ task queue."""
+
+    database_url: str = "postgresql://postgres:postgres@localhost:5433/postgres"
+
+    model_config = {"env_prefix": "PQ_"}

pq/logging.py

@@ -0,0 +1,27 @@
+"""Centralized loguru configuration for PQ."""
+
+import sys
+
+from loguru import logger
+
+
+def configure_logging() -> None:
+    """Configure loguru with a clean, aligned format.
+
+    Called automatically when pq is imported.
+    """
+    # Remove default handler
+    logger.remove()
+
+    # Add handler with clean format
+    # Fixed-width level, simple format without variable-length source info
+    logger.add(
+        sys.stderr,
+        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
+        level="DEBUG",
+        colorize=True,
+    )
+
+
+# Auto-configure on import
+configure_logging()

pq/models.py

@@ -0,0 +1,90 @@
+"""SQLAlchemy 2.0 models for PQ task queue."""
+
+from datetime import datetime, timedelta
+from enum import StrEnum
+from typing import Any
+
+from sqlalchemy import (
+    BigInteger,
+    DateTime,
+    Enum,
+    Identity,
+    Index,
+    Integer,
+    Interval,
+    SmallInteger,
+    String,
+    Text,
+    func,
+)
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+
+class TaskStatus(StrEnum):
+    """Task execution status."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class Base(DeclarativeBase):
+    """Base class for all models."""
+
+    pass
+
+
+class Task(Base):
+    """One-off task with status tracking."""
+
+    __tablename__ = "pq_tasks"
+    __table_args__ = (
+        Index("ix_pq_tasks_status_priority_run_at", "status", "priority", "run_at"),
+    )
+
+    id: Mapped[int] = mapped_column(BigInteger, Identity(), primary_key=True)
+    name: Mapped[str] = mapped_column(String(255), nullable=False)
+    payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False, default=dict)
+    priority: Mapped[int] = mapped_column(SmallInteger, nullable=False, default=0)
+    status: Mapped[TaskStatus] = mapped_column(
+        Enum(TaskStatus, name="task_status", create_constraint=True),
+        nullable=False,
+        default=TaskStatus.PENDING,
+    )
+    run_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), nullable=False, server_default=func.now()
+    )
+    started_at: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True
+    )
+    completed_at: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True
+    )
+    error: Mapped[str | None] = mapped_column(Text, nullable=True)
+    attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+
+
+class Periodic(Base):
+    """Recurring task with interval or cron scheduling."""
+
+    __tablename__ = "pq_periodic"
+    __table_args__ = (
+        Index("ix_pq_periodic_priority_next_run", "priority", "next_run"),
+    )
+
+    id: Mapped[int] = mapped_column(BigInteger, Identity(), primary_key=True)
+    name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
+    payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False, default=dict)
+    priority: Mapped[int] = mapped_column(SmallInteger, nullable=False, default=0)
+    run_every: Mapped[timedelta | None] = mapped_column(Interval, nullable=True)
+    cron: Mapped[str | None] = mapped_column(String(100), nullable=True)
+    next_run: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+    last_run: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True
+    )
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), nullable=False, server_default=func.now()
+    )

pq/priority.py

@@ -0,0 +1,16 @@
+"""Priority levels for task ordering."""
+
+from enum import IntEnum
+
+
+class Priority(IntEnum):
+    """Task priority levels.
+
+    Higher values = higher priority (processed first).
+    """
+
+    BATCH = 0
+    LOW = 25
+    NORMAL = 50
+    HIGH = 75
+    CRITICAL = 100

pq/registry.py

@@ -0,0 +1,74 @@
+"""Function path resolution for task handlers."""
+
+import importlib
+from collections.abc import Callable
+from types import FunctionType
+from typing import Any, cast
+
+
+def get_function_path(func: Callable[..., Any]) -> str:
+    """Get the import path for a function as 'module:name'.
+
+    Args:
+        func: A callable function (must be a top-level function, not a lambda).
+
+    Returns:
+        Import path in 'module:name' format.
+
+    Raises:
+        ValueError: If the function cannot be serialized (lambda, closure, etc).
+    """
+    # Cast to FunctionType for type checker - we verify it has the needed attributes
+    fn = cast(FunctionType, func)
+
+    # Check for lambda
+    if fn.__name__ == "<lambda>":
+        raise ValueError(
+            "Cannot enqueue lambda functions - use a named function instead"
+        )
+
+    # Check for closure (has free variables)
+    if hasattr(fn, "__code__") and fn.__code__.co_freevars:
+        raise ValueError(
+            f"Cannot enqueue closure '{fn.__name__}' - closures capture local variables "
+            "and cannot be imported by the worker"
+        )
+
+    module = getattr(fn, "__module__", None)
+    name = getattr(fn, "__name__", None)
+
+    if module is None or name is None:
+        raise ValueError(f"Cannot determine import path for {func}")
+
+    return f"{module}:{name}"
+
+
+def resolve_function_path(path: str) -> Callable[..., Any]:
+    """Import and return a function from its 'module:name' path.
+
+    Args:
+        path: Import path in 'module:name' format.
+
+    Returns:
+        The imported callable function.
+
+    Raises:
+        ValueError: If the path is invalid or the function cannot be imported.
+    """
+    if ":" not in path:
+        raise ValueError(
+            f"Invalid function path: {path} (expected 'module:name' format)"
+        )
+
+    module_path, func_name = path.rsplit(":", 1)
+
+    try:
+        module = importlib.import_module(module_path)
+        func = getattr(module, func_name)
+        if not callable(func):
+            raise ValueError(f"{path} is not callable")
+        return func
+    except ImportError as e:
+        raise ValueError(f"Cannot import module '{module_path}': {e}") from e
+    except AttributeError as e:
+        raise ValueError(f"Module '{module_path}' has no function '{func_name}'") from e