dwyeapi 0.4.0__py3-none-any.whl

dwyeapi/tasks/__init__.py

@@ -0,0 +1,97 @@
+"""Async task processing system for FastAPI (ARQ-based).
+
+Quick start::
+
+    # 1. Install
+    pip install dwyeapi[tasks]
+
+    # 2. Register tasks
+    from dwyeapi.tasks import register, TaskContext
+
+    @register("process_data")
+    async def process_data(ctx: TaskContext, params: dict):
+        await ctx.log("Processing...")
+        await ctx.update_progress(50)
+        return {"done": True}
+
+    # 3. Mount in FastAPI
+    from dwyeapi.tasks import setup_tasks, task_router
+
+    await setup_tasks(app, settings, session_factory)
+    app.include_router(task_router)
+
+    # 4. Run worker
+    # arq app.worker.WorkerSettings
+
+Public API:
+    setup_tasks     — One-call initialization (ARQ pool + store session factory)
+    task_router     — Ready-to-use APIRouter with CRUD endpoints
+    register        — Decorator to register an async task function
+    TaskContext      — Context object injected into task functions
+    TaskStatus      — Task lifecycle status enum
+    create_worker_settings — Factory to generate ARQ WorkerSettings from BaseSettings
+"""
+
+from __future__ import annotations
+
+from fastapi import FastAPI
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from dwyeapi.config import BaseSettings
+from dwyeapi.tasks.context import TaskContext
+from dwyeapi.tasks.model import TaskStatus
+from dwyeapi.tasks.registry import register
+from dwyeapi.tasks.router import task_router
+from dwyeapi.tasks.worker import create_worker_settings
+
+# Re-export for convenient imports
+__all__ = [
+    "TaskContext",
+    "TaskStatus",
+    "create_worker_settings",
+    "register",
+    "setup_tasks",
+    "task_router",
+]
+
+
+async def setup_tasks(
+    app: FastAPI,
+    settings: BaseSettings,
+    session_factory: async_sessionmaker[AsyncSession],
+) -> None:
+    """Initialize the task system.
+
+    Must be called during application startup (e.g. inside a lifespan).
+
+    1. Configures the ARQ Redis pool from ``settings.redis_url``
+    2. Stores the session_factory for TaskContext database access
+    3. Registers a shutdown hook to close the pool
+
+    Args:
+        app: The FastAPI application instance.
+        settings: An eapi BaseSettings instance with redis_url and task_* fields.
+        session_factory: Async SQLAlchemy session factory for database access.
+
+    Example::
+
+        @asynccontextmanager
+        async def lifespan(app: FastAPI):
+            await setup_tasks(app, settings, session_factory)
+            yield
+
+        app = FastAPI(lifespan=lifespan)
+        app.include_router(task_router)
+    """
+    from dwyeapi.dependencies import create_get_db
+    from dwyeapi.tasks import pool as _pool
+    from dwyeapi.tasks import router as _router
+
+    # Configure and warm up the ARQ pool
+    _pool.configure(settings.redis_url)
+
+    # Wire the db dependency into the router
+    _router._get_db = create_get_db(session_factory)
+
+    # Store session_factory on app state for worker access
+    app.state.task_session_factory = session_factory

dwyeapi/tasks/context.py

@@ -0,0 +1,180 @@
+"""TaskContext — rich context object injected into task functions.
+
+Provides database access, logging, progress tracking, and cancellation
+checking so that business task code stays minimal.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from dwyeapi.tasks import service
+from dwyeapi.tasks.model import TaskStatus
+
+
+class TaskContext:
+    """Execution context passed to every registered task function.
+
+    Provides helpers for logging, progress updates, cancellation checks,
+    and database access so that task functions only contain business logic.
+
+    Attributes:
+        task_id: The unique identifier of the current task.
+        task_type: The registered task type string.
+        params: The parameters submitted with the task.
+    """
+
+    def __init__(
+        self,
+        *,
+        task_id: str,
+        task_type: str,
+        params: dict[str, Any],
+        session_factory: async_sessionmaker[AsyncSession],
+        redis: Any,
+    ) -> None:
+        """Initialize TaskContext.
+
+        Args:
+            task_id: Unique task identifier.
+            task_type: Registered task type name.
+            params: Task input parameters.
+            session_factory: Async SQLAlchemy session factory.
+            redis: ARQ Redis connection for cancel-flag checks.
+        """
+        self.task_id = task_id
+        self.task_type = task_type
+        self.params = params
+        self._session_factory = session_factory
+        self._redis = redis
+        self._cancelled = False
+
+    async def log(self, message: str) -> None:
+        """Append a timestamped log entry.
+
+        Args:
+            message: The log message to record.
+        """
+        async with self._session() as session:
+            await service.append_task_log(session, self.task_id, message)
+
+    async def update_progress(self, value: int, message: str = "") -> None:
+        """Update task progress (0-100) and optionally append a log.
+
+        Args:
+            value: Progress percentage (clamped to 0-100).
+            message: Optional log message to append alongside the progress update.
+        """
+        async with self._session() as session:
+            await service.update_task_progress(session, self.task_id, value)
+            if message:
+                await service.append_task_log(session, self.task_id, message)
+
+    async def is_cancelled(self) -> bool:
+        """Check whether a cancel request has been issued for this task.
+
+        Reads a Redis key ``task_cancel:{task_id}``. Once detected, the
+        result is cached so subsequent calls do not hit Redis again.
+
+        Returns:
+            True if the task has been marked for cancellation.
+        """
+        if self._cancelled:
+            return True
+        val = await self._redis.get(f"task_cancel:{self.task_id}")
+        if val is not None:
+            self._cancelled = True
+        return self._cancelled
+
+    @asynccontextmanager
+    async def db(self) -> AsyncIterator[AsyncSession]:
+        """Provide an async database session via context manager.
+
+        Usage::
+
+            async with ctx.db() as session:
+                result = await session.execute(select(MyModel))
+
+        Yields:
+            An AsyncSession bound to the application's database.
+        """
+        async with self._session_factory() as session:
+            yield session
+
+    @asynccontextmanager
+    async def _session(self) -> AsyncIterator[AsyncSession]:
+        """Internal helper — short-lived session for framework operations."""
+        async with self._session_factory() as session:
+            yield session
+
+
+async def run_task_with_context(
+    task_id: str,
+    task_type: str,
+    params: dict[str, Any],
+    func: Any,
+    session_factory: async_sessionmaker[AsyncSession],
+    redis: Any,
+) -> None:
+    """Execute a task function wrapped in lifecycle management.
+
+    Handles the full task lifecycle:
+    1. Set status to RUNNING
+    2. Create TaskContext and call the business function
+    3. On success → set status to SUCCESS with returned result
+    4. On exception → set status to FAILED with error info
+    5. On cancellation → set status to CANCELED
+
+    Args:
+        task_id: The task's unique identifier.
+        task_type: The registered task type string.
+        params: The task input parameters.
+        func: The async business function to execute.
+        session_factory: Async SQLAlchemy session factory.
+        redis: ARQ Redis connection.
+    """
+    ctx = TaskContext(
+        task_id=task_id,
+        task_type=task_type,
+        params=params,
+        session_factory=session_factory,
+        redis=redis,
+    )
+
+    # Mark as running
+    async with session_factory() as session:
+        await service.update_task_status(session, task_id, TaskStatus.RUNNING)
+        await service.append_task_log(session, task_id, "任务开始执行")
+
+    try:
+        result = await func(ctx, params)
+
+        # Check if cancelled during execution
+        if ctx._cancelled:
+            async with session_factory() as session:
+                await service.update_task_status(session, task_id, TaskStatus.CANCELED)
+                await service.append_task_log(session, task_id, "任务已取消")
+            return
+
+        # Success
+        result_data = result if isinstance(result, dict) else {"result": result}
+        async with session_factory() as session:
+            await service.update_task_status(session, task_id, TaskStatus.SUCCESS, result=result_data)
+            await service.update_task_progress(session, task_id, 100)
+            await service.append_task_log(session, task_id, "任务执行成功")
+
+    except Exception as exc:
+        # Failed
+        async with session_factory() as session:
+            await service.update_task_status(
+                session,
+                task_id,
+                TaskStatus.FAILED,
+                result={"error": str(exc)},
+            )
+            await service.append_task_log(session, task_id, f"任务执行失败: {exc}")
+        raise

dwyeapi/tasks/model.py

@@ -0,0 +1,42 @@
+"""Task ORM model and status enumeration."""
+
+import enum
+import uuid as uuid_lib
+
+from sqlalchemy import JSON, Enum, Integer, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from dwyeapi.database import Base, TimestampMixin
+
+
+class TaskStatus(enum.StrEnum):
+    """Task lifecycle states."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCESS = "success"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+def _generate_task_id() -> str:
+    """Generate a globally unique task ID."""
+    return f"task_{uuid_lib.uuid4().hex}"
+
+
+class Task(Base, TimestampMixin):
+    """Persistent task record.
+
+    Stores task metadata, execution state, progress, logs, and results.
+    Uses the same database as the host application via eapi's Base.
+    """
+
+    __tablename__ = "tasks"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True, default=_generate_task_id)
+    task_type: Mapped[str] = mapped_column(String(50), index=True)
+    status: Mapped[TaskStatus] = mapped_column(Enum(TaskStatus), default=TaskStatus.PENDING)
+    params: Mapped[dict] = mapped_column(JSON)
+    result: Mapped[dict | None] = mapped_column(JSON, nullable=True, default=None)
+    progress: Mapped[int] = mapped_column(Integer, default=0)
+    logs: Mapped[str] = mapped_column(Text, default="")

dwyeapi/tasks/pool.py

@@ -0,0 +1,81 @@
+"""ARQ Redis connection pool management (internal module).
+
+Manages the lifecycle of the ARQ async Redis pool used for enqueuing
+and executing tasks. Not intended for direct use by integrators.
+"""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+from arq import ArqRedis, create_pool
+from arq.connections import RedisSettings
+
+_pool: ArqRedis | None = None
+_redis_settings: RedisSettings | None = None
+
+
+def _parse_redis_url(redis_url: str) -> RedisSettings:
+    """Convert a redis:// URL to ARQ RedisSettings.
+
+    Args:
+        redis_url: Standard Redis connection URL.
+
+    Returns:
+        ARQ RedisSettings instance.
+    """
+    parsed = urlparse(redis_url)
+    return RedisSettings(
+        host=parsed.hostname or "localhost",
+        port=parsed.port or 6379,
+        database=int(parsed.path.lstrip("/") or "0"),
+        password=parsed.password,
+    )
+
+
+def configure(redis_url: str) -> None:
+    """Store Redis settings for later pool creation.
+
+    Args:
+        redis_url: Standard Redis connection URL (e.g. redis://localhost:6379/0).
+    """
+    global _redis_settings
+    _redis_settings = _parse_redis_url(redis_url)
+
+
+def get_redis_settings() -> RedisSettings:
+    """Return the configured ARQ RedisSettings.
+
+    Returns:
+        The RedisSettings instance.
+
+    Raises:
+        RuntimeError: If configure() has not been called.
+    """
+    if _redis_settings is None:
+        msg = "Task pool not configured. Call setup_tasks() first."
+        raise RuntimeError(msg)
+    return _redis_settings
+
+
+async def get_pool() -> ArqRedis:
+    """Get or create the ARQ async Redis pool.
+
+    Returns:
+        The ArqRedis connection pool.
+
+    Raises:
+        RuntimeError: If configure() has not been called.
+    """
+    global _pool
+    if _pool is None:
+        _pool = await create_pool(get_redis_settings())
+    return _pool
+
+
+async def close_pool() -> None:
+    """Close the ARQ pool if it exists. Safe to call multiple times."""
+    global _pool
+    if _pool is not None:
+        await _pool.aclose()
+        _pool = None

dwyeapi/tasks/registry.py

@@ -0,0 +1,83 @@
+"""Task function registry with decorator-based registration."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+# Type alias for an async task function: (TaskContext, dict) -> Any
+TaskFunction = Callable[..., Coroutine[Any, Any, Any]]
+
+
+class TaskRegistry:
+    """Singleton registry that maps task_type strings to async functions.
+
+    Usage::
+
+        from dwyeapi.tasks import register
+
+        @register("process_data")
+        async def process_data(ctx: TaskContext, params: dict):
+            ...
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty registry."""
+        self._tasks: dict[str, TaskFunction] = {}
+
+    def register(self, task_type: str) -> Callable[[TaskFunction], TaskFunction]:
+        """Decorator that registers an async function under a task_type.
+
+        Args:
+            task_type: Unique string identifier for this task.
+
+        Returns:
+            The original function, unmodified.
+
+        Raises:
+            ValueError: If task_type is already registered.
+        """
+
+        def decorator(func: TaskFunction) -> TaskFunction:
+            if task_type in self._tasks:
+                msg = f"Task type '{task_type}' is already registered"
+                raise ValueError(msg)
+            self._tasks[task_type] = func
+            return func
+
+        return decorator
+
+    def get(self, task_type: str) -> TaskFunction | None:
+        """Look up a registered task function by type.
+
+        Args:
+            task_type: The task type to look up.
+
+        Returns:
+            The registered function, or None if not found.
+        """
+        return self._tasks.get(task_type)
+
+    def list_types(self) -> list[str]:
+        """Return all registered task type strings.
+
+        Returns:
+            Sorted list of registered task types.
+        """
+        return sorted(self._tasks.keys())
+
+    def has(self, task_type: str) -> bool:
+        """Check if a task type is registered.
+
+        Args:
+            task_type: The task type to check.
+
+        Returns:
+            True if registered.
+        """
+        return task_type in self._tasks
+
+
+# Module-level singleton
+registry = TaskRegistry()
+register = registry.register

dwyeapi/tasks/router.py

@@ -0,0 +1,196 @@
+"""Ready-to-use FastAPI router for the task system.
+
+Provides endpoints to submit, query, list, and cancel tasks.
+Integrators mount it via ``app.include_router(task_router)``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Callable
+
+from fastapi import APIRouter, Depends, Query
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dwyeapi.exceptions import BusinessError, NotFoundError
+from dwyeapi.response import success
+from dwyeapi.tasks import pool
+from dwyeapi.tasks.model import TaskStatus
+from dwyeapi.tasks.registry import registry
+from dwyeapi.tasks.schema import TaskCreate, TaskListResponse, TaskResponse
+from dwyeapi.tasks.service import append_task_log, create_task, get_task, list_tasks, update_task_progress, update_task_status
+
+task_router = APIRouter(prefix="/tasks", tags=["tasks"])
+
+# This will be set by setup_tasks() at application startup.
+_get_db: Callable[[], AsyncGenerator[AsyncSession]] | None = None
+
+
+async def _db_dependency() -> AsyncGenerator[AsyncSession]:
+    """Async generator that delegates to the actual get_db dependency.
+
+    Yields:
+        An AsyncSession from the configured session factory.
+
+    Raises:
+        RuntimeError: If setup_tasks() has not been called.
+    """
+    if _get_db is None:
+        msg = "Task system not initialized. Call setup_tasks() first."
+        raise RuntimeError(msg)
+    async for session in _get_db():
+        yield session
+
+
+@task_router.post("", summary="提交耗时任务")
+async def submit_task(
+    body: TaskCreate,
+    db: AsyncSession = Depends(_db_dependency),
+) -> dict:
+    """Submit a new async task for background execution.
+
+    Args:
+        body: Task creation payload with task_type and params.
+        db: Injected database session.
+
+    Returns:
+        Unified success response containing the created task.
+    """
+    if not registry.has(body.task_type):
+        raise BusinessError(f"不支持的任务类型: {body.task_type}", code="INVALID_TASK_TYPE")
+
+    task = await create_task(db, body)
+
+    # Enqueue into ARQ
+    arq_pool = await pool.get_pool()
+    await arq_pool.enqueue_job("_task_executor", task.id, body.task_type, body.params)
+
+    return success(data=TaskResponse.model_validate(task).model_dump(mode="json"))
+
+
+@task_router.get("/{task_id}", summary="查询任务状态")
+async def query_task(
+    task_id: str,
+    db: AsyncSession = Depends(_db_dependency),
+) -> dict:
+    """Query the current state of a specific task.
+
+    Args:
+        task_id: The task's unique identifier.
+        db: Injected database session.
+
+    Returns:
+        Unified success response containing the task details.
+    """
+    task = await get_task(db, task_id)
+    if not task:
+        raise NotFoundError("任务")
+    return success(data=TaskResponse.model_validate(task).model_dump(mode="json"))
+
+
+@task_router.get("", summary="任务列表")
+async def query_task_list(
+    db: AsyncSession = Depends(_db_dependency),
+    page: int = Query(default=1, ge=1, description="页码"),
+    page_size: int = Query(default=20, ge=1, le=100, description="每页条数"),
+    status: TaskStatus | None = Query(default=None, description="状态筛选"),
+    task_type: str | None = Query(default=None, description="任务类型筛选"),
+) -> dict:
+    """List tasks with pagination and optional filters.
+
+    Args:
+        db: Injected database session.
+        page: 1-based page number.
+        page_size: Items per page (1-100).
+        status: Optional status filter.
+        task_type: Optional task type filter.
+
+    Returns:
+        Unified success response containing paginated task list.
+    """
+    items, total = await list_tasks(
+        db,
+        page=page,
+        page_size=page_size,
+        status=status,
+        task_type=task_type,
+    )
+    data = TaskListResponse(
+        items=[TaskResponse.model_validate(t) for t in items],
+        total=total,
+    )
+    return success(data=data.model_dump(mode="json"))
+
+
+@task_router.post("/{task_id}/cancel", summary="取消任务")
+async def cancel_task(
+    task_id: str,
+    db: AsyncSession = Depends(_db_dependency),
+) -> dict:
+    """Request cancellation of a running or pending task.
+
+    Sets a Redis flag that the task function can check via ``ctx.is_cancelled()``.
+    Tasks that have already completed cannot be cancelled.
+
+    Args:
+        task_id: The task's unique identifier.
+        db: Injected database session.
+
+    Returns:
+        Unified success response with updated task info.
+    """
+    task = await get_task(db, task_id)
+    if not task:
+        raise NotFoundError("任务")
+
+    terminal_states = {TaskStatus.SUCCESS, TaskStatus.FAILED, TaskStatus.CANCELED}
+    if task.status in terminal_states:
+        raise BusinessError(f"任务已结束, 状态为 {task.status.value}", code="TASK_ALREADY_FINISHED")
+
+    # Set cancel flag in Redis with TTL
+    arq_pool = await pool.get_pool()
+    await arq_pool.set(f"task_cancel:{task_id}", b"1", ex=7200)
+
+    return success(data=TaskResponse.model_validate(task).model_dump(mode="json"))
+
+
+@task_router.post("/{task_id}/retry", summary="重试任务")
+async def retry_task(
+    task_id: str,
+    db: AsyncSession = Depends(_db_dependency),
+) -> dict:
+    """Retry a failed or canceled task.
+
+    Resets status to PENDING, clears result, resets progress, and re-enqueues to ARQ.
+
+    Args:
+        task_id: The task's unique identifier.
+        db: Injected database session.
+
+    Returns:
+        Unified success response with updated task info.
+    """
+    task = await get_task(db, task_id)
+    if not task:
+        raise NotFoundError("任务")
+
+    retryable_states = {TaskStatus.FAILED, TaskStatus.CANCELED}
+    if task.status not in retryable_states:
+        raise BusinessError(f"只有失败或已取消的任务可以重试, 当前状态为 {task.status.value}", code="TASK_NOT_RETRYABLE")
+
+    # Clear cancel flag from Redis
+    arq_pool = await pool.get_pool()
+    await arq_pool.delete(f"task_cancel:{task_id}")
+
+    # Reset task state
+    await update_task_status(db, task_id, TaskStatus.PENDING)
+    await update_task_progress(db, task_id, 0)
+    await append_task_log(db, task_id, "任务重试, 重新入队")
+    task.result = None
+    await db.commit()
+
+    # Re-enqueue
+    await arq_pool.enqueue_job("_task_executor", task.id, task.task_type, task.params)
+
+    # Refresh to get latest state
+    await db.refresh(task)
+    return success(data=TaskResponse.model_validate(task).model_dump(mode="json"))