fix(tenants): preserve RLS context and atomic onboarding

app/config/database.py

@@ -1,4 +1,4 @@
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event, text
 from sqlalchemy.orm import sessionmaker, declarative_base
 from sqlalchemy.exc import SQLAlchemyError
 from fastapi import HTTPException
@@ -31,6 +31,20 @@ engine = create_engine(
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 Base = declarative_base()
 
+@event.listens_for(SessionLocal, "after_begin")
+def restore_request_rls_context(session, transaction, connection):
+    """Reapply request identity after an internal service commit starts a new transaction."""
+    if session.info.get("rls_bypass") is True:
+        connection.execute(text("SELECT set_config('app.bypass_rls', 'on', true)"))
+        return
+
+    tenant_id = session.info.get("rls_tenant_id")
+    if tenant_id:
+        connection.execute(
+            text("SELECT set_config('app.tenant_id', :tenant_id, true)"),
+            {"tenant_id": str(tenant_id)},
+        )
+
 def get_db():
     db = SessionLocal()
     try:
@@ -50,4 +64,4 @@ def get_db():
         db.rollback()
         raise
     finally:
-        db.close()
+        db.close()

app/middleware/auth_middleware.py

@@ -39,6 +39,8 @@ def get_current_user(
         # the database transaction before reading any tenant-owned records.
         token_tenant_id = payload.get("tenant_id")
         if token_tenant_id:
+            db.info.pop("rls_bypass", None)
+            db.info["rls_tenant_id"] = str(token_tenant_id)
             db.execute(
                 text("SELECT set_config('app.tenant_id', :tenant_id, true)"),
                 {"tenant_id": str(token_tenant_id)},
@@ -48,6 +50,8 @@ def get_current_user(
             # Enable transaction-local discovery only long enough to resolve the
             # exact token subject; the result is validated below before platform
             # access remains enabled for the rest of this request.
+            db.info.pop("rls_tenant_id", None)
+            db.info["rls_bypass"] = True
             db.execute(text("SELECT set_config('app.bypass_rls', 'on', true)"))
     except HTTPException:
         raise
@@ -60,6 +64,7 @@ def get_current_user(
     user = db.query(User).filter(User.id == user_id).first()
     if user is None:
         if not token_tenant_id:
+            db.info.pop("rls_bypass", None)
             db.execute(text("SELECT set_config('app.bypass_rls', 'off', true)"))
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
@@ -75,6 +80,7 @@ def get_current_user(
     elif user.tenant_id is not None:
         # Fail closed for a validly signed tenant-user token that is malformed,
         # stale, or missing its canonical tenant claim.
+        db.info.pop("rls_bypass", None)
         db.execute(text("SELECT set_config('app.bypass_rls', 'off', true)"))
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,

app/services/auth/role_service.py

@@ -19,7 +19,12 @@ logger = logging.getLogger(__name__)
 class RoleService:
 
     @staticmethod
-    def create_role(db: Session, role_data: RoleCreate, emit_events: bool = True) -> Role:
+    def create_role(
+        db: Session,
+        role_data: RoleCreate,
+        emit_events: bool = True,
+        commit: bool = True,
+    ) -> Role:
         existing = (
             db.query(Role)
             .filter(
@@ -42,11 +47,16 @@ class RoleService:
         )
 
         db.add(role)
-        db.commit()
-        db.refresh(role)
+        if commit:
+            db.commit()
+            db.refresh(role)
+        else:
+            db.flush()
 
         if role_data.access_ids:
-            RoleService.assign_accesses(db, role.id, role_data.access_ids)
+            RoleService.assign_accesses(
+                db, role.id, role_data.access_ids, commit=commit
+            )
             
         if emit_events:
             assigned_modules = (
@@ -104,12 +114,20 @@ class RoleService:
                         tenant_id=role.tenant_id
                     )
 
-                    db.commit()
+                    if commit:
+                        db.commit()
+                    else:
+                        db.flush()
 
         return role
 
     @staticmethod
-    def assign_accesses(db: Session, role_id: uuid.UUID, access_ids: List[uuid.UUID]):
+    def assign_accesses(
+        db: Session,
+        role_id: uuid.UUID,
+        access_ids: List[uuid.UUID],
+        commit: bool = True,
+    ):
         db.query(RoleAccess).filter(RoleAccess.role_id == role_id).delete()
         db.query(RoleModuleAccess).filter(RoleModuleAccess.role_id == role_id).delete()
 
@@ -129,7 +147,10 @@ class RoleService:
             for access in module_accesses:
                 db.add(RoleModuleAccess(role_id=role_id, module_access_id=access.id))
 
-        db.commit()
+        if commit:
+            db.commit()
+        else:
+            db.flush()
 
     @staticmethod
     def get_all_roles(db: Session, tenant_id: uuid.UUID = None):
@@ -361,4 +382,4 @@ class RoleService:
             page=page,
             page_size=page_size,
             total_pages=total_pages,
-        )
+        )

app/services/auth/tenant_service.py

@@ -162,7 +162,12 @@ class TenantService:
                 is_default=True,
                 access_ids=all_access_ids
             )
-            role = RoleService.create_role(db, role_create_data, emit_events=False)
+            role = RoleService.create_role(
+                db,
+                role_create_data,
+                emit_events=False,
+                commit=False,
+            )
 
             role_module_perms = (
                 db.query(RoleModuleAccess, ModuleAccess)