Update code generation and task return types

Author: wu-clan

backend/app/generator/api/v1/gen.py

@@ -8,7 +8,7 @@
 from backend.app.generator.conf import generator_settings
 from backend.app.generator.schema.gen import ImportParam
 from backend.app.generator.service.gen_service import gen_service
-from backend.common.response.response_schema import ResponseModel, response_base
+from backend.common.response.response_schema import ResponseModel, ResponseSchemaModel, response_base
 from backend.common.security.jwt import DependsJwtAuth
 from backend.common.security.permission import RequestPermission
 from backend.common.security.rbac import DependsRBAC
@@ -17,7 +17,9 @@
 
 
 @router.get('/tables', summary='获取数据库表')
-async def get_all_tables(table_schema: Annotated[str, Query(..., description='数据库名')] = 'fba') -> ResponseModel:
+async def get_all_tables(
+    table_schema: Annotated[str, Query(..., description='数据库名')] = 'fba',
+) -> ResponseSchemaModel[list[str]]:
     data = await gen_service.get_tables(table_schema=table_schema)
     return response_base.success(data=data)
 
@@ -36,13 +38,13 @@ async def import_table(obj: ImportParam) -> ResponseModel:
 
 
 @router.get('/preview/{pk}', summary='生成代码预览', dependencies=[DependsJwtAuth])
-async def preview_code(pk: Annotated[int, Path(..., description='业务ID')]) -> ResponseModel:
+async def preview_code(pk: Annotated[int, Path(..., description='业务ID')]) -> ResponseSchemaModel[dict[str, bytes]]:
     data = await gen_service.preview(pk=pk)
     return response_base.success(data=data)
 
 
 @router.get('/generate/{pk}/path', summary='获取代码生成路径', dependencies=[DependsJwtAuth])
-async def generate_path(pk: Annotated[int, Path(..., description='业务ID')]):
+async def generate_path(pk: Annotated[int, Path(..., description='业务ID')]) -> ResponseSchemaModel[list[str]]:
     data = await gen_service.get_generate_path(pk=pk)
     return response_base.success(data=data)
 

backend/app/generator/api/v1/gen_business.py

@@ -4,14 +4,15 @@
 
 from fastapi import APIRouter, Depends, Path
 
+from backend.app.generator.model import GenBusiness, GenModel
 from backend.app.generator.schema.gen_business import (
     CreateGenBusinessParam,
     GetGenBusinessDetail,
     UpdateGenBusinessParam,
 )
 from backend.app.generator.service.gen_business_service import gen_business_service
 from backend.app.generator.service.gen_model_service import gen_model_service
-from backend.common.response.response_schema import ResponseModel, response_base
+from backend.common.response.response_schema import ResponseModel, ResponseSchemaModel, response_base
 from backend.common.security.jwt import DependsJwtAuth
 from backend.common.security.permission import RequestPermission
 from backend.common.security.rbac import DependsRBAC
@@ -21,21 +22,21 @@
 
 
 @router.get('/all', summary='获取所有代码生成业务', dependencies=[DependsJwtAuth])
-async def get_all_businesses() -> ResponseModel:
+async def get_all_businesses() -> ResponseSchemaModel[list[GenBusiness]]:
     businesses = await gen_business_service.get_all()
     data = select_list_serialize(businesses)
     return response_base.success(data=data)
 
 
 @router.get('/{pk}', summary='获取代码生成业务详情', dependencies=[DependsJwtAuth])
-async def get_business(pk: Annotated[int, Path(...)]) -> ResponseModel:
+async def get_business(pk: Annotated[int, Path(...)]) -> ResponseSchemaModel[GetGenBusinessDetail]:
     business = await gen_business_service.get(pk=pk)
     data = GetGenBusinessDetail(**select_as_dict(business))
     return response_base.success(data=data)
 
 
 @router.get('/{pk}/models', summary='获取代码生成业务所有模型', dependencies=[DependsJwtAuth])
-async def get_business_all_models(pk: Annotated[int, Path(...)]) -> ResponseModel:
+async def get_business_all_models(pk: Annotated[int, Path(...)]) -> ResponseSchemaModel[list[GenModel]]:
     models = await gen_model_service.get_by_business(business_id=pk)
     data = select_list_serialize(models)
     return response_base.success(data=data)

backend/app/generator/api/v1/gen_model.py

@@ -6,7 +6,7 @@
 
 from backend.app.generator.schema.gen_model import CreateGenModelParam, GetGenModelDetail, UpdateGenModelParam
 from backend.app.generator.service.gen_model_service import gen_model_service
-from backend.common.response.response_schema import ResponseModel, response_base
+from backend.common.response.response_schema import ResponseModel, ResponseSchemaModel, response_base
 from backend.common.security.jwt import DependsJwtAuth
 from backend.common.security.permission import RequestPermission
 from backend.common.security.rbac import DependsRBAC
@@ -16,13 +16,13 @@
 
 
 @router.get('/types', summary='获取代码生成模型列类型', dependencies=[DependsJwtAuth])
-async def get_model_types() -> ResponseModel:
+async def get_model_types() -> ResponseSchemaModel[list[str]]:
     model_types = await gen_model_service.get_types()
     return response_base.success(data=model_types)
 
 
 @router.get('/{pk}', summary='获取代码生成模型详情', dependencies=[DependsJwtAuth])
-async def get_model(pk: Annotated[int, Path(...)]) -> ResponseModel:
+async def get_model(pk: Annotated[int, Path(...)]) -> ResponseSchemaModel[GetGenModelDetail]:
     model = await gen_model_service.get(pk=pk)
     data = GetGenModelDetail(**select_as_dict(model))
     return response_base.success(data=data)

backend/app/generator/service/gen_service.py

@@ -69,7 +69,7 @@ async def import_business_and_model(*, obj: ImportParam) -> None:
                 await gen_model_dao.create(db, CreateGenModelParam(**model_data), pd_type=pd_type)
 
     @staticmethod
-    async def render_tpl_code(*, business: GenBusiness) -> dict:
+    async def render_tpl_code(*, business: GenBusiness) -> dict[str, str]:
         gen_models = await gen_model_service.get_by_business(business_id=business.id)
         if not gen_models:
             raise errors.NotFoundError(msg='代码生成模型表为空')
@@ -79,7 +79,7 @@ async def render_tpl_code(*, business: GenBusiness) -> dict:
             tpl_code_map[tpl_path] = await gen_template.get_template(tpl_path).render_async(**gen_vars)
         return tpl_code_map
 
-    async def preview(self, *, pk: int) -> dict:
+    async def preview(self, *, pk: int) -> dict[str, bytes]:
         async with async_db_session() as db:
             business = await gen_business_dao.get(db, pk)
             if not business:
@@ -91,7 +91,7 @@ async def preview(self, *, pk: int) -> dict:
             }
 
     @staticmethod
-    async def get_generate_path(*, pk: int) -> list:
+    async def get_generate_path(*, pk: int) -> list[str]:
         async with async_db_session() as db:
             business = await gen_business_dao.get(db, pk)
             if not business:

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

@@ -4,9 +4,9 @@
 
 from fastapi import APIRouter, Depends, Path
 
-from backend.app.task.schema.task import RunParam
+from backend.app.task.schema.task import RunParam, TaskResult
 from backend.app.task.service.task_service import task_service
-from backend.common.response.response_schema import ResponseModel, response_base
+from backend.common.response.response_schema import ResponseModel, ResponseSchemaModel, response_base
 from backend.common.security.jwt import DependsJwtAuth
 from backend.common.security.permission import RequestPermission
 from backend.common.security.rbac import DependsRBAC
@@ -15,7 +15,7 @@
 
 
 @router.get('', summary='获取可执行任务', dependencies=[DependsJwtAuth])
-async def get_all_tasks() -> ResponseModel:
+async def get_all_tasks() -> ResponseSchemaModel[list[str]]:
     tasks = await task_service.get_list()
     return response_base.success(data=tasks)
 
@@ -27,7 +27,7 @@ async def get_all_tasks() -> ResponseModel:
     description='此接口被视为作废，建议使用 flower 查看任务详情',
     dependencies=[DependsJwtAuth],
 )
-async def get_task_detail(tid: Annotated[str, Path(description='任务ID')]) -> ResponseModel:
+async def get_task_detail(tid: Annotated[str, Path(description='任务ID')]) -> ResponseSchemaModel[TaskResult]:
     status = task_service.get_detail(tid=tid)
     return response_base.success(data=status)
 
@@ -53,6 +53,6 @@ async def revoke_task(tid: Annotated[str, Path(description='任务ID')]) -> Resp
         DependsRBAC,
     ],
 )
-async def run_task(obj: RunParam) -> ResponseModel:
+async def run_task(obj: RunParam) -> ResponseSchemaModel[str]:
     task = task_service.run(obj=obj)
     return response_base.success(data=task)

backend/app/task/schema/task.py

@@ -9,3 +9,15 @@ class RunParam(SchemaBase):
     name: str = Field(description='任务名称')
     args: list | None = Field(default=None, description='任务函数位置参数')
     kwargs: dict | 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

backend/app/task/service/task_service.py

@@ -1,25 +1,26 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-
 from celery.exceptions import NotRegistered
 from celery.result import AsyncResult
 from starlette.concurrency import run_in_threadpool
 
 from backend.app.task.celery import celery_app
-from backend.app.task.schema.task import RunParam
-from backend.common.dataclasses import TaskResult
+from backend.app.task.schema.task import RunParam, TaskResult
+from backend.common.exception import errors
 from backend.common.exception.errors import NotFoundError
 
 
 class TaskService:
     @staticmethod
-    async def get_list():
+    async def get_list() -> list[str]:
         registered_tasks = await run_in_threadpool(celery_app.control.inspect().registered)
+        if not registered_tasks:
+            raise errors.ForbiddenError(msg='celery 服务未启动')
         tasks = list(registered_tasks.values())[0]
         return tasks
 
     @staticmethod
-    def get_detail(*, tid: str):
+    def get_detail(*, tid: str) -> TaskResult:
         try:
             result = AsyncResult(id=tid, app=celery_app)
         except NotRegistered:
@@ -45,7 +46,7 @@ def revoke(*, tid: str):
         result.revoke(terminate=True)
 
     @staticmethod
-    def run(*, obj: RunParam):
+    def run(*, obj: RunParam) -> str:
         task: AsyncResult = celery_app.send_task(name=obj.name, args=obj.args, kwargs=obj.kwargs)
         return task.task_id
 

backend/common/dataclasses.py

@@ -52,16 +52,3 @@ class AccessToken:
 class RefreshToken:
     refresh_token: str
     refresh_token_expire_time: datetime
-
-
-@dataclasses.dataclass
-class TaskResult:
-    result: str
-    traceback: str
-    status: str
-    name: str
-    args: list | None
-    kwargs: dict | None
-    worker: str
-    retries: int | None
-    queue: str | None