轮询任务调度

Author: EITS403

app/tasks/sqlmap_scheduler.py

@@ -5,9 +5,18 @@
 from app.tasks.sqlmap_worker import poll_single_sqlmap_task
 
 
-@shared_task
-def poll_active_sqlmap_tasks():
-    with SessionLocal() as session:
+@shared_task(
+    bind=True,
+    autoretry_for=(Exception,),
+    retry_backoff=10,
+    retry_kwargs={"max_retries": 5},
+)
+def poll_active_sqlmap_tasks(self):
+    """
+    轮询所有 pending / running 的 sqlmap 任务
+    """
+    session = SessionLocal()
+    try:
         tasks = (
             session.query(SqlmapScanPayload)
             .filter(
@@ -18,3 +27,6 @@ def poll_active_sqlmap_tasks():
 
         for task in tasks:
             poll_single_sqlmap_task.delay(task.task_id)
+
+    finally:
+        session.close()