53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
"""Celery tasks for data collection"""
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
from typing import Dict, Any
|
|
|
|
from app.db.session import async_session_factory
|
|
from app.services.collectors.registry import collector_registry
|
|
|
|
|
|
async def run_collector_task(collector_name: str) -> Dict[str, Any]:
|
|
"""Run a single collector task"""
|
|
collector = collector_registry.get(collector_name)
|
|
if not collector:
|
|
return {"status": "failed", "error": f"Collector {collector_name} not found"}
|
|
|
|
if not collector_registry.is_active(collector_name):
|
|
return {"status": "skipped", "reason": "Collector is disabled"}
|
|
|
|
async with async_session_factory() as db:
|
|
from app.models.task import CollectionTask
|
|
from app.models.datasource import DataSource
|
|
|
|
# Find datasource
|
|
result = await db.execute(
|
|
"SELECT id FROM data_sources WHERE collector_class = :class_name",
|
|
{"class_name": f"{collector.__class__.__name__}"},
|
|
)
|
|
datasource = result.fetchone()
|
|
|
|
task = CollectionTask(
|
|
datasource_id=datasource[0] if datasource else 0,
|
|
status="running",
|
|
started_at=datetime.utcnow(),
|
|
)
|
|
db.add(task)
|
|
await db.commit()
|
|
|
|
result = await collector.run(db)
|
|
|
|
task.status = result["status"]
|
|
task.completed_at = datetime.utcnow()
|
|
task.records_processed = result.get("records_processed", 0)
|
|
task.error_message = result.get("error")
|
|
await db.commit()
|
|
|
|
return result
|
|
|
|
|
|
def run_collector_sync(collector_name: str) -> Dict[str, Any]:
|
|
"""Synchronous wrapper for running collectors"""
|
|
return asyncio.run(run_collector_task(collector_name))
|