Optimize task

Author: wu-clan

backend/app/task/api/router.py

@@ -5,6 +5,6 @@
 from backend.app.task.api.v1.task import router as task_router
 from backend.core.conf import settings
 
-v1 = APIRouter(prefix=settings.FASTAPI_API_V1_PATH)
+v1 = APIRouter(prefix=settings.FASTAPI_API_V1_PATH, tags=['任务'])
 
-v1.include_router(task_router, prefix='/tasks', tags=['任务'])
+v1.include_router(task_router, prefix='/tasks')

backend/app/task/api/v1/task.py

@@ -27,7 +27,7 @@ async def get_all_tasks() -> ResponseSchemaModel[list[str]]:
     description='此接口被视为作废，建议使用 flower 查看任务详情',
     dependencies=[DependsJwtAuth],
 )
-async def get_task_detail(tid: Annotated[str, Path(description='任务ID')]) -> ResponseSchemaModel[TaskResult]:
+async def get_task_detail(tid: Annotated[str, Path(description='任务 UUID')]) -> ResponseSchemaModel[TaskResult]:
     status = task_service.get_detail(tid=tid)
     return response_base.success(data=status)
 
@@ -40,7 +40,7 @@ async def get_task_detail(tid: Annotated[str, Path(description='任务ID')]) ->
         DependsRBAC,
     ],
 )
-async def revoke_task(tid: Annotated[str, Path(description='任务ID')]) -> ResponseModel:
+async def revoke_task(tid: Annotated[str, Path(description='任务 UUID')]) -> ResponseModel:
     task_service.revoke(tid=tid)
     return response_base.success()
 

backend/app/task/celery.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from typing import Any
+
 import celery
 import celery_aio_pool
 
@@ -9,61 +11,64 @@
 __all__ = ['celery_app']
 
 
-def init_celery() -> celery.Celery:
-    """初始化 celery 应用"""
-
-    # TODO: Update this work if celery version >= 6.0.0
-    # https://github.com/fastapi-practices/fastapi_best_architecture/issues/321
-    # https://github.com/celery/celery/issues/7874
-    celery.app.trace.build_tracer = celery_aio_pool.build_async_tracer
-    celery.app.trace.reset_worker_optimizations()
-
-    # Celery Schedule Tasks
-    # https://docs.celeryq.dev/en/stable/userguide/periodic-tasks.html
-    beat_schedule = task_settings.CELERY_SCHEDULE
-
-    # Celery Config
-    # https://docs.celeryq.dev/en/stable/userguide/configuration.html
-    broker_url = (
-        (
+def get_broker_url() -> str:
+    """获取消息代理 URL"""
+    if task_settings.CELERY_BROKER == 'redis':
+        return (
             f'redis://:{settings.REDIS_PASSWORD}@{settings.REDIS_HOST}:'
             f'{settings.REDIS_PORT}/{task_settings.CELERY_BROKER_REDIS_DATABASE}'
         )
-        if task_settings.CELERY_BROKER == 'redis'
-        else (
-            f'amqp://{task_settings.RABBITMQ_USERNAME}:{task_settings.RABBITMQ_PASSWORD}@'
-            f'{task_settings.RABBITMQ_HOST}:{task_settings.RABBITMQ_PORT}'
-        )
+    return (
+        f'amqp://{task_settings.RABBITMQ_USERNAME}:{task_settings.RABBITMQ_PASSWORD}@'
+        f'{task_settings.RABBITMQ_HOST}:{task_settings.RABBITMQ_PORT}'
     )
-    result_backend = (
+
+
+def get_result_backend() -> str:
+    """获取结果后端 URL"""
+    return (
         f'redis://:{settings.REDIS_PASSWORD}@{settings.REDIS_HOST}:'
         f'{settings.REDIS_PORT}/{task_settings.CELERY_BACKEND_REDIS_DATABASE}'
     )
-    result_backend_transport_options = {
-        'global_keyprefix': f'{task_settings.CELERY_BACKEND_REDIS_PREFIX}',
+
+
+def get_result_backend_transport_options() -> dict[str, Any]:
+    """获取结果后端传输选项"""
+    return {
+        'global_keyprefix': task_settings.CELERY_BACKEND_REDIS_PREFIX,
         'retry_policy': {
             'timeout': task_settings.CELERY_BACKEND_REDIS_TIMEOUT,
         },
     }
 
+
+def init_celery() -> celery.Celery:
+    """初始化 Celery 应用"""
+
+    # TODO: Update this work if celery version >= 6.0.0
+    # https://github.com/fastapi-practices/fastapi_best_architecture/issues/321
+    # https://github.com/celery/celery/issues/7874
+    celery.app.trace.build_tracer = celery_aio_pool.build_async_tracer
+    celery.app.trace.reset_worker_optimizations()
+
     app = celery.Celery(
         'fba_celery',
         enable_utc=False,
         timezone=settings.DATETIME_TIMEZONE,
-        beat_schedule=beat_schedule,
-        broker_url=broker_url,
+        beat_schedule=task_settings.CELERY_SCHEDULE,
+        broker_url=get_broker_url(),
         broker_connection_retry_on_startup=True,
-        result_backend=result_backend,
-        result_backend_transport_options=result_backend_transport_options,
+        result_backend=get_result_backend(),
+        result_backend_transport_options=get_result_backend_transport_options(),
         task_cls='app.task.celery_task.base:TaskBase',
         task_track_started=True,
     )
 
-    # Load task modules
+    # 自动发现任务模块
     app.autodiscover_tasks(task_settings.CELERY_TASK_PACKAGES)
 
     return app
 
 
-# 创建 celery 实例
+# 创建 Celery 实例
 celery_app: celery.Celery = init_celery()

backend/app/task/celery_task/base.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from typing import Any
 
 from celery import Task
 from sqlalchemy.exc import SQLAlchemyError
@@ -9,16 +10,37 @@
 
 
 class TaskBase(Task):
-    """任务基类"""
+    """Celery 任务基类"""
 
     autoretry_for = (SQLAlchemyError,)
     max_retries = task_settings.CELERY_TASK_MAX_RETRIES
 
-    async def before_start(self, task_id, args, kwargs):
+    async def before_start(self, task_id: str, args, kwargs) -> None:
+        """
+        任务开始前执行钩子
+
+        :param task_id: 任务 ID
+        :return:
+        """
         await task_notification(msg=f'任务 {task_id} 开始执行')
 
-    async def on_success(self, retval, task_id, args, kwargs):
+    async def on_success(self, retval: Any, task_id: str, args, kwargs) -> None:
+        """
+        任务成功后执行钩子
+
+        :param retval: 任务返回值
+        :param task_id: 任务 ID
+        :return:
+        """
         await task_notification(msg=f'任务 {task_id} 执行成功')
 
-    async def on_failure(self, exc, task_id, args, kwargs, einfo):
+    async def on_failure(self, exc: Exception, task_id: str, args, kwargs, einfo) -> None:
+        """
+        任务失败后执行钩子
+
+        :param exc: 异常对象
+        :param task_id: 任务 ID
+        :param einfo: 异常信息
+        :return:
+        """
         await task_notification(msg=f'任务 {task_id} 执行失败')

backend/app/task/celery_task/tasks.py

@@ -7,5 +7,6 @@
 
 @celery_app.task(name='task_demo_async')
 async def task_demo_async() -> str:
+    """异步示例任务，模拟耗时操作"""
     await sleep(20)
     return 'test async'

backend/app/task/conf.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 from functools import lru_cache
-from typing import Literal
+from typing import Any, Literal
 
 from celery.schedules import crontab
 from pydantic import model_validator
@@ -11,25 +11,22 @@
 
 
 class TaskSettings(BaseSettings):
-    """Task Settings"""
+    """Celery 任务配置"""
 
     model_config = SettingsConfigDict(env_file=f'{BASE_PATH}/.env', env_file_encoding='utf-8', extra='ignore')
 
-    # Env Config
-    ENVIRONMENT: Literal['dev', 'pro']
-
-    # Env Celery
-    CELERY_BROKER_REDIS_DATABASE: int  # 仅在 dev 模式时生效
+    # .env Redis 配置
+    CELERY_BROKER_REDIS_DATABASE: int
     CELERY_BACKEND_REDIS_DATABASE: int
 
-    # Env Rabbitmq
+    # .env RabbitMQ 配置
     # docker run -d --hostname fba-mq --name fba-mq  -p 5672:5672 -p 15672:15672 rabbitmq:latest
     RABBITMQ_HOST: str
     RABBITMQ_PORT: int
     RABBITMQ_USERNAME: str
     RABBITMQ_PASSWORD: str
 
-    # Celery
+    # Celery 基础配置
     CELERY_BROKER: Literal['rabbitmq', 'redis'] = 'redis'
     CELERY_BACKEND_REDIS_PREFIX: str = 'fba:celery:'
     CELERY_BACKEND_REDIS_TIMEOUT: int = 5
@@ -38,7 +35,9 @@ class TaskSettings(BaseSettings):
         'app.task.celery_task.db_log',
     ]
     CELERY_TASK_MAX_RETRIES: int = 5
-    CELERY_SCHEDULE: dict = {
+
+    # Celery 定时任务配置
+    CELERY_SCHEDULE: dict[str, dict[str, Any]] = {
         'exec-every-10-seconds': {
             'task': 'task_demo_async',
             'schedule': 10,
@@ -55,15 +54,16 @@ class TaskSettings(BaseSettings):
 
     @model_validator(mode='before')
     @classmethod
-    def validate_celery_broker(cls, values):
+    def validate_celery_broker(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """生产环境强制使用 RabbitMQ 作为消息代理"""
         if values['ENVIRONMENT'] == 'pro':
             values['CELERY_BROKER'] = 'rabbitmq'
         return values
 
 
 @lru_cache
 def get_task_settings() -> TaskSettings:
-    """获取 task 配置"""
+    """获取 Celery 任务配置"""
     return TaskSettings()
 
 

backend/app/task/schema/task.py

@@ -1,23 +1,29 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from typing import Any
+
 from pydantic import Field
 
 from backend.common.schema import SchemaBase
 
 
 class RunParam(SchemaBase):
+    """任务运行参数"""
+
     name: str = Field(description='任务名称')
-    args: list | None = Field(default=None, description='任务函数位置参数')
-    kwargs: dict | None = Field(default=None, description='任务函数关键字参数')
+    args: list[Any] | None = Field(default=None, description='任务函数位置参数')
+    kwargs: dict[str, Any] | None = Field(default=None, description='任务函数关键字参数')
 
 
 class TaskResult(SchemaBase):
-    result: str
-    traceback: str
-    status: str
-    name: str
-    args: list | None
-    kwargs: dict | None
-    worker: str
-    retries: int | None
-    queue: str | None
+    """任务执行结果"""
+
+    result: str = Field(description='任务执行结果')
+    traceback: str = Field(description='错误堆栈信息')
+    status: str = Field(description='任务状态')
+    name: str = Field(description='任务名称')
+    args: list[Any] | None = Field(default=None, description='任务函数位置参数')
+    kwargs: dict[str, Any] | None = Field(default=None, description='任务函数关键字参数')
+    worker: str = Field(description='执行任务的 worker')
+    retries: int | None = Field(default=None, description='重试次数')
+    queue: str | None = Field(default=None, description='任务队列')

backend/app/task/service/task_service.py

@@ -13,14 +13,21 @@
 class TaskService:
     @staticmethod
     async def get_list() -> list[str]:
+        """获取所有已注册的 Celery 任务列表"""
         registered_tasks = await run_in_threadpool(celery_app.control.inspect().registered)
         if not registered_tasks:
-            raise errors.ForbiddenError(msg='celery 服务未启动')
+            raise errors.ForbiddenError(msg='Celery 服务未启动')
         tasks = list(registered_tasks.values())[0]
         return tasks
 
     @staticmethod
     def get_detail(*, tid: str) -> TaskResult:
+        """
+        获取指定任务的详细信息
+
+        :param tid: 任务 UUID
+        :return:
+        """
         try:
             result = AsyncResult(id=tid, app=celery_app)
         except NotRegistered:
@@ -38,7 +45,13 @@ def get_detail(*, tid: str) -> TaskResult:
         )
 
     @staticmethod
-    def revoke(*, tid: str):
+    def revoke(*, tid: str) -> None:
+        """
+        撤销指定的任务
+
+        :param tid: 任务 UUID
+        :return:
+        """
         try:
             result = AsyncResult(id=tid, app=celery_app)
         except NotRegistered:
@@ -47,6 +60,12 @@ def revoke(*, tid: str):
 
     @staticmethod
     def run(*, obj: RunParam) -> str:
+        """
+        运行指定的任务
+
+        :param obj: 任务运行参数
+        :return:
+        """
         task: AsyncResult = celery_app.send_task(name=obj.name, args=obj.args, kwargs=obj.kwargs)
         return task.task_id