添加任务自动注册

Author: EITS403

app/middleware/celery_app.py

@@ -22,6 +22,8 @@
     backend=f"redis://:{db_password}@{redis_host}:{redis_port}/{redis_db}",
 )
 
+celery_app.autodiscover_tasks(packages=["app.tasks"])
+
 celery_app.conf.update(
     task_serializer="json",
     accept_content=["json"],

app/middleware/celery_beat.py

@@ -1,5 +1,4 @@
-from celery.schedules import crontab
-from app.middleware import celery_app
+from app.middleware.celery_app import celery_app
 
 celery_app.conf.beat_schedule = {
     "poll-sqlmap-tasks-every-5-seconds": {

app/tasks/__init__.py

@@ -0,0 +1,2 @@
+from . import sqlmap_scheduler
+from . import sqlmap_worker

app/tasks/sqlmap_scheduler.py

@@ -1,11 +1,10 @@
-from sqlalchemy import select
-from app.middleware.celery_app import celery_app
+from celery import shared_task
 from app.database.celery_sync_database import SessionLocal
 from app.models.sqlmap_result import SqlmapScanPayload, ScanStatus
 from app.tasks.sqlmap_worker import poll_single_sqlmap_task
 
 
-@celery_app.task(name="app.tasks.sqlmap_scheduler.poll_active_sqlmap_tasks")
+@shared_task
 def poll_active_sqlmap_tasks():
     with SessionLocal() as session:
         tasks = (

app/tasks/sqlmap_worker.py

@@ -1,12 +1,9 @@
 import requests
 from datetime import datetime
-from sqlalchemy.exc import SQLAlchemyError
-
-from app.middleware.celery_app import celery_app
+from celery import shared_task
 from app.database.celery_sync_database import SessionLocal
 from app.models.sqlmap_result import (
     SqlmapScanPayload,
-    SqlmapScanLog,
     ScanStatus,
     SqlmapScanResult,
 )
@@ -57,12 +54,11 @@ def normalize_sqlmap_result(raw: dict) -> dict:
     return result
 
 
-@celery_app.task(
+@shared_task(
     bind=True,
     autoretry_for=(Exception,),
     retry_backoff=5,
     retry_kwargs={"max_retries": 3},
-    name="app.tasks.sqlmap_worker.poll_single_sqlmap_task",
 )
 def poll_single_sqlmap_task(self, task_id: str):
     session = SessionLocal()