dev

coldStartTasks/ai_pipeline/basic.py

@@ -3,8 +3,10 @@ import time
 import datetime
 import requests
 
+from applications.db import DatabaseConnector
 
-def get_status_field_by_process(process):
+
+def get_status_field_by_process(process: str) -> tuple[str, str]:
     match process:
         case "upload":
             status = "upload_status"
@@ -24,7 +26,11 @@ def get_status_field_by_process(process):
 
 
 def roll_back_lock_tasks(
-    db_client, process, max_process_time, init_status, processing_status
+    db_client: DatabaseConnector,
+    process: str,
+    max_process_time: int,
+    init_status: int,
+    processing_status: int
 ) -> int:
     """
     rollback tasks which have been locked for a long time
@@ -74,10 +80,13 @@ def generate_summary_prompt(text):
     return prompt
 
 
-def update_task_queue_status(db_client, task_id, process, ori_status, new_status):
-    """
-    回滚长时间处于处理中的任务
-    """
+def update_task_queue_status(
+        db_client: DatabaseConnector,
+        task_id: int,
+        process: str,
+        ori_status: int,
+        new_status: int) -> int:
+    # update task queue status
     status, update_timestamp = get_status_field_by_process(process)
     update_query = f"""
         update video_content_understanding 
@@ -96,15 +105,11 @@ def update_task_queue_status(db_client, task_id, process, ori_status, new_status
     return roll_back_rows
 
 
-def update_video_pool_status(db_client, content_trace_id, ori_status, new_status):
-    """
-    回滚长时间处于处理中的任务
-    """
-    update_sql = f"""
-        update publish_single_video_source
-        set status = %s
-        where content_trace_id = %s and status = %s;
-    """
+def update_video_pool_status(
+        db_client: DatabaseConnector,
+        content_trace_id: str,
+        ori_status: int,
+        new_status: int) -> int:
     # update publish_single_source_status
     update_query = f"""
                 update publish_single_video_source