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 = os.getenv("VALKEY_URL", "redis://valkey:6379/0") 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( "anexo76_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", ""), task_id, headers, ) logger.info( "Celery task_prerun RLS context task=%s task_id=%s tenant_id=%s company_id=%s", getattr(task, "name", ""), 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", ""), task_id, rls_tenant_var.get(), rls_company_var.get(), ) reset_rls_context_tokens(token_t, token_c) delattr(task, _RLS_TOKENS_ATTR) # ---------------------------------------------------------------------------- # Import models in correct order for SQLAlchemy relationship resolution # MUST happen AFTER celery_app exists to avoid circular imports during # initialization when models trigger route/task imports. # ---------------------------------------------------------------------------- # Orden: PedimentoCode y RegimenPedimento antes de CodePedimentoRegimen (mapper) from api.v1.modules.public.reference_data.pedimento_codes.models import PedimentoCode from api.v1.modules.public.reference_data.pedimento_regimens.models import RegimenPedimento from api.v1.modules.public.reference_data.code_pedimento_regimens.models import ( CodePedimentoRegimen, ) # InvoiceType debe cargarse antes de InvoiceHeader (FK invoice_header.invoice_type -> public.invoice_types.key) from api.v1.modules.public.reference_data.invoice_types.models import InvoiceType # noqa: F401 # CustomsSection debe cargarse antes de InvoiceComplianceMx (FK invoice_compliance_mx.aduana -> public.customs_sections.customs_code) from api.v1.modules.public.reference_data.customs_sections.models import CustomsSection # noqa: F401 # CRITICAL: FaLineItem must be imported BEFORE LineItem from api.v1.modules.a24.fa.fa_item_lines.models import FaLineItem # noqa: F401 from api.v1.modules.a76.items.models import LineItem # noqa: F401 # CRITICAL: BalanceMovement must be loaded before DischargeDetail (FK a24.balance_movement) from api.v1.modules.a24.balance_movements.models import BalanceMovement # noqa: F401 from api.v1.modules.a24.discharges.models import DischargeHeader, DischargeDetail # noqa: F401 celery_app.conf.update( include=[ "api.v1.modules.a76.reports.importacion.facturas.task", "api.v1.modules.a76.reports.importacion.consolidados.task", "api.v1.modules.a76.reports.importacion.packing_list.task", "api.v1.modules.a76.reports.exportacion.aviso_consolidado.task", "api.v1.modules.a76.reports.movements.invoices.tasks", "api.v1.modules.a76.reports.movements.saldos.tasks", "api.v1.modules.a76.reports.exportacion.descargo.task", "api.v1.modules.a76.layouts_csv.facturas.tasks", "api.v1.modules.a76.layouts_csv.exportacion.tasks", "api.v1.modules.a76.layouts_csv.cambio_regimen_regularizacion.tasks", "api.v1.modules.a76.layouts_csv.customs_brokers.tasks", "api.v1.modules.a76.layouts_csv.clients_and_providers.tasks", "api.v1.modules.a76.layouts_csv.pedmientos.tasks", "api.v1.modules.a76.layouts_csv.exchange_rate.tasks", "api.v1.modules.a76.layouts_csv.us_tariff_fractions.tasks", "api.v1.modules.a76.layouts_csv.classes.tasks", "api.v1.modules.a76.layouts_csv.parts.tasks", "api.v1.modules.a76.layouts_csv.boms.tasks", "api.v1.modules.a76.layouts_csv.vehicles.tasks", "api.v1.modules.a76.layouts_csv.drivers.tasks", "api.v1.modules.a76.layouts_csv.trailers.tasks", "api.v1.modules.a76.layouts_csv.transportistas.tasks", "api.v1.modules.a76.reports.exportacion.transmission.MAINX30.task", "api.v1.modules.a76.reports.importacion.transmission.temporal.MAINX30.task", "api.v1.modules.a76.reports.importacion.transmission.definitive.MAINX30.task", "api.v1.modules.a76.reports.importacion.winsaai.invoices.task", "api.v1.modules.a76.reports.importacion.winsaai.pedimentos.task", "api.v1.modules.core.help_center.tasks", "api.v1.modules.core.help_center.tasks", "api.v1.modules.a76.invoices.imports.process.task", "api.v1.modules.a76.invoices.imports.revert.task", "api.v1.modules.a76.invoices.exports.process.task", "api.v1.modules.a76.invoices.exports.revert.task", "api.v1.modules.a76.layouts_csv.common.victor", "api.v1.modules.a76.factura_cove.tasks", "api.v1.modules.a76.expediente_archivos.tasks", ] # Ruta al módulo donde están las tareas ) # 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()