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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user