feat: add Celery task infrastructure for module system
Phase 4 of module migration plan: - Add ScheduledTask dataclass for declaring Celery Beat tasks - Add tasks_path and scheduled_tasks fields to ModuleDefinition - Create ModuleTask base class with database session management - Create task discovery utilities (discover_module_tasks, build_beat_schedule) - Update celery_config.py to discover and register module tasks - Maintain backward compatibility with legacy task modules Modules can now define tasks in their tasks/ directory and scheduled tasks in their definition. The infrastructure supports gradual migration of existing tasks from app/tasks/ to their respective modules. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
172
app/modules/task_base.py
Normal file
172
app/modules/task_base.py
Normal file
@@ -0,0 +1,172 @@
|
||||
# app/modules/task_base.py
|
||||
"""
|
||||
Base Celery task class for module tasks.
|
||||
|
||||
Provides a ModuleTask base class that handles:
|
||||
- Database session lifecycle (create/close)
|
||||
- Context manager pattern for session usage
|
||||
- Proper cleanup on task completion or failure
|
||||
- Structured logging with task context
|
||||
|
||||
This is the standard base class for all Celery tasks defined within modules.
|
||||
It replaces the legacy app/tasks/celery_tasks/base.py for new module tasks.
|
||||
|
||||
Usage:
|
||||
from app.core.celery_config import celery_app
|
||||
from app.modules.task_base import ModuleTask
|
||||
|
||||
@celery_app.task(bind=True, base=ModuleTask)
|
||||
def my_task(self, arg1, arg2):
|
||||
with self.get_db() as db:
|
||||
# Use db session
|
||||
result = db.query(Model).filter(...).all()
|
||||
return result
|
||||
"""
|
||||
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
from celery import Task
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import SessionLocal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModuleTask(Task):
|
||||
"""
|
||||
Base Celery task with database session management.
|
||||
|
||||
Provides:
|
||||
- Database session context manager
|
||||
- Automatic session cleanup on success/failure
|
||||
- Structured logging for task lifecycle events
|
||||
- Retry tracking
|
||||
|
||||
All module tasks should use this as their base class for consistent
|
||||
behavior and proper resource management.
|
||||
|
||||
Example:
|
||||
@celery_app.task(bind=True, base=ModuleTask)
|
||||
def process_subscription(self, subscription_id: int):
|
||||
with self.get_db() as db:
|
||||
subscription = db.query(Subscription).get(subscription_id)
|
||||
# Process subscription...
|
||||
return {"status": "processed"}
|
||||
"""
|
||||
|
||||
# Mark as abstract so Celery doesn't register this as a task
|
||||
abstract = True
|
||||
|
||||
@contextmanager
|
||||
def get_db(self) -> Session:
|
||||
"""
|
||||
Context manager for database session.
|
||||
|
||||
Yields a database session and ensures proper cleanup on both
|
||||
success and failure. Commits on success, rolls back on error.
|
||||
|
||||
Yields:
|
||||
Session: SQLAlchemy database session
|
||||
|
||||
Example:
|
||||
with self.get_db() as db:
|
||||
vendor = db.query(Vendor).filter(Vendor.id == vendor_id).first()
|
||||
vendor.status = "active"
|
||||
# Session commits automatically on exit
|
||||
"""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Database error in task {self.name}: {e}",
|
||||
extra={
|
||||
"task_name": self.name,
|
||||
"task_id": getattr(self.request, "id", None),
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def on_failure(
|
||||
self,
|
||||
exc: Exception,
|
||||
task_id: str,
|
||||
args: tuple,
|
||||
kwargs: dict,
|
||||
einfo: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Called when task fails.
|
||||
|
||||
Logs the failure with structured context for debugging and monitoring.
|
||||
"""
|
||||
logger.error(
|
||||
f"Task {self.name}[{task_id}] failed: {exc}",
|
||||
extra={
|
||||
"task_name": self.name,
|
||||
"task_id": task_id,
|
||||
"args": args,
|
||||
"kwargs": kwargs,
|
||||
"exception": str(exc),
|
||||
"traceback": str(einfo),
|
||||
},
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def on_success(
|
||||
self,
|
||||
retval: Any,
|
||||
task_id: str,
|
||||
args: tuple,
|
||||
kwargs: dict,
|
||||
) -> None:
|
||||
"""
|
||||
Called when task succeeds.
|
||||
|
||||
Logs successful completion with task context.
|
||||
"""
|
||||
logger.info(
|
||||
f"Task {self.name}[{task_id}] completed successfully",
|
||||
extra={
|
||||
"task_name": self.name,
|
||||
"task_id": task_id,
|
||||
},
|
||||
)
|
||||
|
||||
def on_retry(
|
||||
self,
|
||||
exc: Exception,
|
||||
task_id: str,
|
||||
args: tuple,
|
||||
kwargs: dict,
|
||||
einfo: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Called when task is being retried.
|
||||
|
||||
Logs retry attempt with reason and retry count.
|
||||
"""
|
||||
retry_count = getattr(self.request, "retries", 0)
|
||||
logger.warning(
|
||||
f"Task {self.name}[{task_id}] retrying (attempt {retry_count + 1}): {exc}",
|
||||
extra={
|
||||
"task_name": self.name,
|
||||
"task_id": task_id,
|
||||
"retry_count": retry_count,
|
||||
"exception": str(exc),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Alias for backward compatibility and clarity
|
||||
DatabaseTask = ModuleTask
|
||||
|
||||
__all__ = ["ModuleTask", "DatabaseTask"]
|
||||
Reference in New Issue
Block a user