添加异步和同步添加数据库函数

Author: EITS403

app/core/sqlmap_core.py

@@ -1,26 +1,31 @@
 from sqlalchemy import select
 
 from app.database.database import AsyncSessionLocal
+from app.middleware.custom_decorators import with_async_session
 from app.models.sqlmap_result import SqlmapScanPayload, SqlmapScanLog
+from app.database.celery_sync_database import SessionLocal
 
 
+# 初次创建任务后将会存储数据库
+@with_async_session
 async def task_add(
-    task_id: int,
+    *,
+    session,
+    task_id: str,
     scan_url: str,
     status: str,
     scan_risk: int = 1,
     scan_level: int = 1,
 ):
-    async with AsyncSessionLocal() as session:
-        task = SqlmapScanPayload(
-            task_id=task_id,
-            scan_url=scan_url,
-            status=status,
-            scan_risk=scan_risk,
-            scan_level=scan_level,
-        )
-        session.add(task)
-        await session.commit()
+    task = SqlmapScanPayload(
+        task_id=task_id,
+        scan_url=scan_url,
+        status=status,
+        scan_risk=scan_risk,
+        scan_level=scan_level,
+    )
+    session.add(task)
+    await session.commit()
 
 
 async def list_tasks():
@@ -45,3 +50,24 @@ async def get_task_logs(task_id: str, limit: int = 100, offset: int = 0):
         )
         logs = result.scalars().all()
     return logs
+
+
+# 同步扫描任务写入。防止数据库丢失
+def celery_task_add(
+    *,
+    session,
+    task_id: str,
+    scan_url: str,
+    status: str,
+    scan_risk: int = 1,
+    scan_level: int = 1,
+):
+    task = SqlmapScanPayload(
+        task_id=task_id,
+        scan_url=scan_url,
+        status=status,
+        scan_risk=scan_risk,
+        scan_level=scan_level,
+    )
+    session.add(task)
+    session.commit()