Implement row-level security (RLS) context management for database sessions. Refactor invoice processing and reverting tasks to utilize scoped database sessions with RLS context. Update middleware to extract and set company ID from requests. Enhance task dispatching to propagate RLS context via Celery headers. Update architecture documentation to reflect RLS implementation details.

This commit is contained in:
2026-04-24 18:01:46 -05:00
parent 7acf77994f
commit 73a6d92863
13 changed files with 1052 additions and 306 deletions

View File

@@ -1,5 +1,8 @@
import os
from celery import Celery
from celery.signals import task_postrun, task_prerun
from core.database import rls_company_var, rls_tenant_var
valkey_url = os.getenv("VALKEY_URL", "redis://valkey:6379/0")
@@ -13,6 +16,59 @@ celery_app = Celery(
)
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"))
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
try:
rls_tenant_var.reset(token_t)
rls_company_var.reset(token_c)
except ValueError:
rls_tenant_var.set(None)
rls_company_var.set(None)
finally:
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