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.

2026-04-24 18:01:46 -05:00
parent 7acf77994f
commit 73a6d92863
13 changed files with 1052 additions and 306 deletions
--- a/backend/core/database.py
+++ b/backend/core/database.py
@@ -2,13 +2,21 @@
 Configuración de base de datos con soporte multi-tenant
 - Base de datos compartida (core_db) para tenants pequeños/medianos
 - Bases de datos dedicadas para clientes enterprise
+
+Row-Level Security (RLS):
+    Para respetar el aislamiento por tenant/company definido en PostgreSQL,
+    cada sesión fija las GUCs ``app.tenant_id`` y ``app.company_id`` vía
+    ``SET LOCAL`` al inicio de cada transacción. El listener
+    ``after_begin`` aplica el contexto guardado en ``Session.info``.
 """

 import logging
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from typing import AsyncGenerator, Dict, Generator, Optional

-from sqlalchemy import create_engine
+from fastapi import Request
+from sqlalchemy import create_engine, event, text
 from sqlalchemy.exc import ProgrammingError
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 from sqlalchemy.orm import Session, declarative_base, sessionmaker
@@ -17,10 +25,8 @@ from .config import settings

 logger = logging.getLogger(__name__)

-# Base declarativa para modelos ORM
 Base = declarative_base()

-# Engine y SessionLocal para base de datos core (sincrónico)
 core_engine = create_engine(
    settings.core_database_url,
    pool_pre_ping=True,
@@ -32,7 +38,6 @@ core_engine = create_engine(
 CoreSessionLocal = sessionmaker(
    autocommit=False, autoflush=False, bind=core_engine)

-# Engine asíncrono para operaciones async
 async_core_engine = create_async_engine(
    settings.async_core_database_url,
    pool_pre_ping=True,
@@ -45,27 +50,144 @@ AsyncCoreSessionLocal = async_sessionmaker(
    async_core_engine, class_=AsyncSession, expire_on_commit=False
 )

-# Cache de engines para tenants con BD dedicada
 _tenant_engines: Dict[str, any] = {}


-def get_core_db() -> Generator[Session, None, None]:
+RLS_TENANT_KEY = "rls_tenant_id"
+RLS_COMPANY_KEY = "rls_company_id"
+
+rls_tenant_var: ContextVar[Optional[int]] = ContextVar("rls_tenant_id", default=None)
+rls_company_var: ContextVar[Optional[int]] = ContextVar("rls_company_id", default=None)
+
+
+def _apply_rls_context(connection, tenant_id: Optional[int], company_id: Optional[int]) -> None:
+    """Ejecuta ``SET LOCAL`` en la transacción activa para fijar el contexto RLS."""
+    tenant_value = "" if tenant_id is None else str(int(tenant_id))
+    company_value = "" if company_id is None else str(int(company_id))
+    connection.execute(
+        text(
+            "SELECT set_config('app.tenant_id', :t, true), "
+            "set_config('app.company_id', :c, true)"
+        ),
+        {"t": tenant_value, "c": company_value},
+    )
+
+
+def _resolve_context(session) -> tuple[Optional[int], Optional[int]]:
+    """Selecciona tenant_id/company_id desde ``session.info`` y, si faltan,
+    desde las ContextVars (usadas por tareas Celery vía task_prerun)."""
+    tenant_id = session.info.get(RLS_TENANT_KEY)
+    company_id = session.info.get(RLS_COMPANY_KEY)
+    if tenant_id is None:
+        tenant_id = rls_tenant_var.get()
+    if company_id is None:
+        company_id = rls_company_var.get()
+    return tenant_id, company_id
+
+
+@event.listens_for(Session, "after_begin")
+def _after_begin(session: Session, transaction, connection) -> None:  # type: ignore[no-untyped-def]
+    """Aplica ``SET LOCAL`` en cada transacción nueva.
+
+    Cubre sesiones síncronas y asíncronas porque ``AsyncSession`` envuelve
+    internamente una ``Session`` que hereda de esta clase.
    """
-    Dependency para obtener sesión de base de datos core (compartida)
-    Uso en FastAPI: db: Session = Depends(get_core_db)
+    tenant_id, company_id = _resolve_context(session)
+    if tenant_id is None and company_id is None:
+        return
+    _apply_rls_context(connection, tenant_id, company_id)
+
+
+def set_rls_context(
+    session: Session,
+    tenant_id: Optional[int] = None,
+    company_id: Optional[int] = None,
+) -> None:
+    """Guarda el contexto RLS en la sesión y, si hay transacción abierta, lo aplica.
+
+    Útil para endpoints que validan acceso a una compañía específica después de
+    crear la sesión (por ejemplo rutas que reciben ``company_id`` en el path).
    """
+    session.info[RLS_TENANT_KEY] = tenant_id
+    session.info[RLS_COMPANY_KEY] = company_id
+    if session.in_transaction():
+        _apply_rls_context(session.connection(), tenant_id, company_id)
+
+
+def _extract_rls_context(request: Optional[Request]) -> tuple[Optional[int], Optional[int]]:
+    """Recupera ``tenant_id`` / ``company_id`` del estado del request (o de cookies)."""
+    if request is None:
+        return None, None
+    tenant_id = getattr(request.state, "tenant_id", None)
+    company_id = getattr(request.state, "company_id", None)
+    if company_id is None:
+        cookie_value = request.cookies.get("active_company_id")
+        if cookie_value:
+            try:
+                company_id = int(cookie_value)
+            except (TypeError, ValueError):
+                company_id = None
+    return tenant_id, company_id
+
+
+def get_core_db(request: Request = None) -> Generator[Session, None, None]:
+    """Dependency para obtener sesión síncrona con contexto RLS.
+
+    FastAPI inyecta ``Request`` automáticamente; los llamadores existentes que
+    escriben ``db: Session = Depends(get_core_db)`` siguen funcionando sin
+    cambios porque ``Request`` se resuelve en la capa de dependencia.
+    """
+    tenant_id, company_id = _extract_rls_context(request)
    db = CoreSessionLocal()
+    db.info[RLS_TENANT_KEY] = tenant_id
+    db.info[RLS_COMPANY_KEY] = company_id
    try:
        yield db
    finally:
        db.close()


-async def get_async_core_db() -> AsyncGenerator[AsyncSession, None]:
-    """
-    Dependency para obtener sesión asíncrona de base de datos core
-    """
+async def get_async_core_db(request: Request = None) -> AsyncGenerator[AsyncSession, None]:
+    """Dependency async para obtener sesión con contexto RLS."""
+    tenant_id, company_id = _extract_rls_context(request)
    async with AsyncCoreSessionLocal() as session:
+        session.info[RLS_TENANT_KEY] = tenant_id
+        session.info[RLS_COMPANY_KEY] = company_id
+        try:
+            yield session
+        finally:
+            await session.close()
+
+
+@contextmanager
+def scoped_core_db(
+    tenant_id: Optional[int] = None,
+    company_id: Optional[int] = None,
+) -> Generator[Session, None, None]:
+    """Abre una sesión síncrona con contexto RLS explícito.
+
+    Pensado para tareas Celery, comandos de mantenimiento o cualquier camino
+    fuera del ciclo de request HTTP. El contexto se aplica con ``SET LOCAL``
+    en cada transacción.
+    """
+    db = CoreSessionLocal()
+    db.info[RLS_TENANT_KEY] = tenant_id
+    db.info[RLS_COMPANY_KEY] = company_id
+    try:
+        yield db
+    finally:
+        db.close()
+
+
+@asynccontextmanager
+async def scoped_async_core_db(
+    tenant_id: Optional[int] = None,
+    company_id: Optional[int] = None,
+) -> AsyncGenerator[AsyncSession, None]:
+    """Variante async de :func:`scoped_core_db`."""
+    async with AsyncCoreSessionLocal() as session:
+        session.info[RLS_TENANT_KEY] = tenant_id
+        session.info[RLS_COMPANY_KEY] = company_id
        try:
            yield session
        finally:
@@ -73,16 +195,7 @@ async def get_async_core_db() -> AsyncGenerator[AsyncSession, None]:


 def get_tenant_engine(tenant_id: int, db_config: dict):
-    """
-    Obtiene o crea un engine para un tenant con BD dedicada
-
-    Args:
-        tenant_id: ID del tenant
-        db_config: Configuración de BD {host, port, name, user, password}
-
-    Returns:
-        Engine de SQLAlchemy para el tenant
-    """
+    """Obtiene o crea un engine para un tenant con BD dedicada."""
    if tenant_id not in _tenant_engines:
        db_url = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['name']}"
        _tenant_engines[tenant_id] = create_engine(
@@ -95,21 +208,16 @@ def get_tenant_engine(tenant_id: int, db_config: dict):
 def get_tenant_db(
    tenant_id: int, db_config: Optional[dict] = None
 ) -> Generator[Session, None, None]:
-    """
-    Context manager para obtener sesión de BD de un tenant específico
+    """Context manager para obtener sesión de BD de un tenant específico.

-    Si db_config es None, usa la BD core (compartida)
-    Si db_config está presente, usa la BD dedicada del tenant
-
-    Uso:
-        with get_tenant_db(tenant_id, config) as db:
-            # operaciones con db
+    Si ``db_config`` es ``None`` usa la BD core compartida y aplica RLS con
+    el ``tenant_id`` recibido. Si el tenant tiene BD dedicada, el aislamiento
+    es físico y no se fija contexto RLS (no hay columna ``tenant_id``).
    """
    if db_config is None:
-        # Tenant en BD compartida
        db = CoreSessionLocal()
+        db.info[RLS_TENANT_KEY] = tenant_id
    else:
-        # Tenant con BD dedicada
        engine = get_tenant_engine(tenant_id, db_config)
        SessionLocal = sessionmaker(
            autocommit=False, autoflush=False, bind=engine)
@@ -122,13 +230,10 @@ def get_tenant_db(


 def init_db():
-    """
-    Inicializa las tablas de la base de datos core
-    """
+    """Inicializa las tablas de la base de datos core."""
    try:
        Base.metadata.create_all(bind=core_engine, checkfirst=True)
    except ProgrammingError as e:
-        # Si la tabla ya existe, es seguro continuar
        if "already exists" in str(e):
            logger.warning(
                f"Algunas tablas ya existen en la base de datos: {e}")
@@ -137,8 +242,6 @@ def init_db():


 async def init_async_db():
-    """
-    Inicializa las tablas de la base de datos core (async)
-    """
+    """Inicializa las tablas de la base de datos core (async)."""
    async with async_core_engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)