项目/AI生成链路合并

video-gen-api/app/services/module_generation_flow_base_service.py

@@ -16,6 +16,11 @@ from app.models.module_generation_step import ModuleGenerationStep
 from app.models.user import User
 from app.enums.module_generation_flow import ModuleGenerationFlowConfig
 from app.services.module_generation_step_common_service import build_step_input, build_step_output, utc_now
+from app.services.generation.pipeline.db_lock_service import (
+    apply_short_lock_timeout,
+    execute_with_lock_timeout,
+    raise_if_database_lock_busy,
+)
 from app.services.resource_accounting_service import (
     SOURCE_MODEL_CHAT_TASK,
     soft_delete_resources_by_source,
@@ -44,8 +49,13 @@ async def get_project_for_user(
     if populate_existing:
         query = query.execution_options(populate_existing=True)
     if for_update:
+        await apply_short_lock_timeout(db)
         query = query.with_for_update()
-    result = await db.execute(query.limit(1))
+    try:
+        result = await db.execute(query.limit(1))
+    except Exception as exc:
+        raise_if_database_lock_busy(exc)
+        raise
     project = result.scalar_one_or_none()
     if not project:
         raise HTTPException(status_code=404, detail=config.project_not_found_message)
@@ -72,8 +82,13 @@ async def get_step_for_user(
     if not user.is_admin:
         query = query.where(ModuleGenerationStep.user_id == user.id)
     if for_update:
+        await apply_short_lock_timeout(db)
         query = query.with_for_update()
-    result = await db.execute(query.limit(1))
+    try:
+        result = await db.execute(query.limit(1))
+    except Exception as exc:
+        raise_if_database_lock_busy(exc)
+        raise
     step = result.scalar_one_or_none()
     if not step:
         raise HTTPException(status_code=404, detail=config.step_not_found_message)
@@ -278,7 +293,9 @@ async def load_chat_tasks_for_steps(
     )
     if for_update:
         stmt = stmt.with_for_update()
-    result = await db.execute(stmt)
+        result = await execute_with_lock_timeout(db, stmt)
+    else:
+        result = await db.execute(stmt)
     return {task.id: task for task in result.scalars().all()}
 
 
@@ -314,7 +331,8 @@ async def assert_project_has_no_active_chat_tasks(
     config: ModuleGenerationFlowConfig,
     detail_message: str = "当前存在生成中任务，请等待生成完成或失败后再操作",
 ) -> dict[str, ChatGenerationTask]:
-    result = await db.execute(
+    result = await execute_with_lock_timeout(
+        db,
         select(ModuleGenerationStep)
         .where(
             ModuleGenerationStep.project_id == project.id,
@@ -353,7 +371,8 @@ async def soft_delete_steps_from_index(
     - 已完成任务只做软删任务与 generated_resources，释放容量统计；失败任务只软删任务。
     """
     deleted_at = deleted_at or utc_now()
-    result = await db.execute(
+    result = await execute_with_lock_timeout(
+        db,
         select(ModuleGenerationStep)
         .where(
             ModuleGenerationStep.project_id == project.id,