127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
import os
|
|
import logging
|
|
from celery import Celery
|
|
from celery.signals import task_postrun, task_prerun
|
|
|
|
from core.database import reset_rls_context_tokens, rls_company_var, rls_tenant_var
|
|
from core.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
valkey_url = settings.VALKEY_URL
|
|
print(f"DEBUG: Celery Broker URL: {valkey_url}")
|
|
logger.info(
|
|
"Initializing Celery app app_version=%s environment=%s broker=%s",
|
|
settings.APP_VERSION,
|
|
settings.ENVIRONMENT,
|
|
valkey_url,
|
|
)
|
|
|
|
# Configurar broker y backend explícitamente en el constructor
|
|
celery_app = Celery(
|
|
"app_tasks",
|
|
broker=valkey_url,
|
|
backend=valkey_url,
|
|
)
|
|
celery_app.set_default()
|
|
|
|
|
|
_RLS_TOKENS_ATTR = "_rls_context_tokens"
|
|
|
|
|
|
def _coerce_int(value) -> int | None:
|
|
if value is None or value == "":
|
|
return None
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
@task_prerun.connect
|
|
def _set_rls_context_from_task(task_id=None, task=None, args=None, kwargs=None, **_):
|
|
"""Fija las ContextVars de RLS para la ejecución de la tarea.
|
|
|
|
Las rutas propagan ``tenant_id`` / ``company_id`` vía Celery headers en
|
|
:func:`track_and_dispatch`. Aquí los materializamos en ContextVars para
|
|
que cualquier sesión que se abra durante la tarea (incluidos los helpers
|
|
``scoped_core_db`` y llamadas directas a ``CoreSessionLocal()``) aplique
|
|
``SET LOCAL`` automáticamente.
|
|
"""
|
|
headers = {}
|
|
request = getattr(task, "request", None) if task is not None else None
|
|
if request is not None:
|
|
headers = getattr(request, "headers", None) or {}
|
|
|
|
tenant_id = _coerce_int(headers.get("rls_tenant_id"))
|
|
company_id = _coerce_int(headers.get("rls_company_id"))
|
|
if tenant_id is None:
|
|
logger.warning(
|
|
"Celery task_prerun missing rls_tenant_id task=%s task_id=%s headers=%s",
|
|
getattr(task, "name", "<unknown>"),
|
|
task_id,
|
|
headers,
|
|
)
|
|
logger.info(
|
|
"Celery task_prerun RLS context task=%s task_id=%s tenant_id=%s company_id=%s",
|
|
getattr(task, "name", "<unknown>"),
|
|
task_id,
|
|
tenant_id,
|
|
company_id,
|
|
)
|
|
|
|
token_t = rls_tenant_var.set(tenant_id)
|
|
token_c = rls_company_var.set(company_id)
|
|
setattr(task, _RLS_TOKENS_ATTR, (token_t, token_c))
|
|
|
|
|
|
@task_postrun.connect
|
|
def _reset_rls_context_from_task(task_id=None, task=None, **_):
|
|
"""Restaura las ContextVars al terminar la tarea (evita fuga entre tareas
|
|
cuando un worker reutiliza el mismo hilo)."""
|
|
tokens = getattr(task, _RLS_TOKENS_ATTR, None) if task is not None else None
|
|
if tokens is None:
|
|
return
|
|
token_t, token_c = tokens
|
|
logger.info(
|
|
"Celery task_postrun clearing RLS context task=%s task_id=%s tenant_id=%s company_id=%s",
|
|
getattr(task, "name", "<unknown>"),
|
|
task_id,
|
|
rls_tenant_var.get(),
|
|
rls_company_var.get(),
|
|
)
|
|
reset_rls_context_tokens(token_t, token_c)
|
|
delattr(task, _RLS_TOKENS_ATTR)
|
|
|
|
celery_app.conf.update(
|
|
include=[
|
|
"api.v1.modules.core.help_center.tasks",
|
|
# Agrega aquí las tareas de tu proyecto:
|
|
# "api.v1.modules.example.tasks",
|
|
]
|
|
)
|
|
|
|
# Configuraciones adicionales
|
|
celery_app.conf.update(
|
|
task_track_started=True,
|
|
task_serializer="json",
|
|
accept_content=["json"],
|
|
result_serializer="json",
|
|
timezone="America/Mexico_City",
|
|
enable_utc=True,
|
|
)
|
|
|
|
celery_app.conf.beat_schedule = {
|
|
"sync-from-hub-every-minute": {
|
|
"task": "sync_from_hub_task",
|
|
"schedule": 60.0, # Run every 60 seconds
|
|
},
|
|
"cleanup-orphan-layout-imports-hourly": {
|
|
"task": "cleanup_orphan_layout_imports",
|
|
"schedule": 3600.0,
|
|
},
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
celery_app.start()
|