fix(rls): Apply persistent RLS transactions (#7917)

Co-authored-by: Víctor Fernández Poyatos <victor@prowler.com>

api/CHANGELOG.md

@@ -24,7 +24,7 @@ All notable changes to the **Prowler API** are documented in this file.
 ## [v1.8.1] (Prowler v5.7.1)
 
 ### Fixed
-- Added database index to improve performance on finding lookup. [(#7800)](https://github.com/prowler-cloud/prowler/pull/7800)
+- Added database index to improve performance on finding lookup [(#7800)](https://github.com/prowler-cloud/prowler/pull/7800).
 
 ---
 

api/pyproject.toml

@@ -35,7 +35,7 @@ name = "prowler-api"
 package-mode = false
 # Needed for the SDK compatibility
 requires-python = ">=3.11,<3.13"
-version = "1.8.2"
+version = "1.8.3"
 
 [project.scripts]
 celery = "src.backend.config.settings.celery"

api/src/backend/api/base_views.py

@@ -47,11 +47,9 @@ class BaseViewSet(ModelViewSet):
 
 
 class BaseRLSViewSet(BaseViewSet):
-    def dispatch(self, request, *args, **kwargs):
-        with transaction.atomic():
-            return super().dispatch(request, *args, **kwargs)
-
     def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+
         # Ideally, this logic would be in the `.setup()` method but DRF view sets don't call it
         # https://docs.djangoproject.com/en/5.1/ref/class-based-views/base/#django.views.generic.base.View.setup
         if request.auth is None:
@@ -61,9 +59,19 @@ class BaseRLSViewSet(BaseViewSet):
         if tenant_id is None:
             raise NotAuthenticated("Tenant ID is not present in token")
 
-        with rls_transaction(tenant_id):
-            self.request.tenant_id = tenant_id
-            return super().initial(request, *args, **kwargs)
+        self.request.tenant_id = tenant_id
+
+        self._rls_cm = rls_transaction(tenant_id)
+        self._rls_cm.__enter__()
+
+    def finalize_response(self, request, response, *args, **kwargs):
+        response = super().finalize_response(request, response, *args, **kwargs)
+
+        if hasattr(self, "_rls_cm"):
+            self._rls_cm.__exit__(None, None, None)
+            del self._rls_cm
+
+        return response
 
     def get_serializer_context(self):
         context = super().get_serializer_context()
@@ -109,6 +117,8 @@ class BaseTenantViewset(BaseViewSet):
                 pass  # Tenant might not exist, handle gracefully
 
     def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+
         if request.auth is None:
             raise NotAuthenticated
 
@@ -117,19 +127,27 @@ class BaseTenantViewset(BaseViewSet):
             raise NotAuthenticated("Tenant ID is not present in token")
 
         user_id = str(request.user.id)
-        with rls_transaction(value=user_id, parameter=POSTGRES_USER_VAR):
-            return super().initial(request, *args, **kwargs)
+
+        self._rls_cm = rls_transaction(value=user_id, parameter=POSTGRES_USER_VAR)
+        self._rls_cm.__enter__()
+
+    def finalize_response(self, request, response, *args, **kwargs):
+        response = super().finalize_response(request, response, *args, **kwargs)
+
+        if hasattr(self, "_rls_cm"):
+            self._rls_cm.__exit__(None, None, None)
+            del self._rls_cm
+
+        return response
 
 
 class BaseUserViewset(BaseViewSet):
-    def dispatch(self, request, *args, **kwargs):
-        with transaction.atomic():
-            return super().dispatch(request, *args, **kwargs)
-
     def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+
         # TODO refactor after improving RLS on users
         if request.stream is not None and request.stream.method == "POST":
-            return super().initial(request, *args, **kwargs)
+            return
         if request.auth is None:
             raise NotAuthenticated
 
@@ -137,6 +155,16 @@ class BaseUserViewset(BaseViewSet):
         if tenant_id is None:
             raise NotAuthenticated("Tenant ID is not present in token")
 
-        with rls_transaction(tenant_id):
-            self.request.tenant_id = tenant_id
-            return super().initial(request, *args, **kwargs)
+        self.request.tenant_id = tenant_id
+
+        self._rls_cm = rls_transaction(tenant_id)
+        self._rls_cm.__enter__()
+
+    def finalize_response(self, request, response, *args, **kwargs):
+        response = super().finalize_response(request, response, *args, **kwargs)
+
+        if hasattr(self, "_rls_cm"):
+            self._rls_cm.__exit__(None, None, None)
+            del self._rls_cm
+
+        return response

api/src/backend/api/specs/v1.yaml

@@ -1,7 +1,7 @@
 openapi: 3.0.3
 info:
   title: Prowler API
-  version: 1.8.2
+  version: 1.8.3
   description: |-
     Prowler API specification.
 

api/src/backend/api/v1/views.py

@@ -260,7 +260,7 @@ class SchemaView(SpectacularAPIView):
 
     def get(self, request, *args, **kwargs):
         spectacular_settings.TITLE = "Prowler API"
-        spectacular_settings.VERSION = "1.8.2"
+        spectacular_settings.VERSION = "1.8.3"
         spectacular_settings.DESCRIPTION = (
             "Prowler API specification.\n\nThis file is auto-generated."
         )