From 18301997da7790e32264b7e55d334c3eb15d39a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Fern=C3=A1ndez=20Poyatos?=
 <victor@prowler.com>
Date: Tue, 3 Jun 2025 16:38:00 +0200
Subject: [PATCH] revert(views): calling order to initial view method (#7921)

(cherry picked from commit 6e7a32cb516e813b441e0fe2f7766a9ce270a488)

# Conflicts:
#	api/src/backend/api/base_views.py
---
 api/src/backend/api/base_views.py | 55 ++++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/api/src/backend/api/base_views.py b/api/src/backend/api/base_views.py
index 605afe332d..f8ed8f7c47 100644
--- a/api/src/backend/api/base_views.py
+++ b/api/src/backend/api/base_views.py
@@ -52,6 +52,8 @@ class BaseRLSViewSet(BaseViewSet):
             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 +63,25 @@ class BaseRLSViewSet(BaseViewSet):
         if tenant_id is None:
             raise NotAuthenticated("Tenant ID is not present in token")
 
+<<<<<<< HEAD
         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
+>>>>>>> 6e7a32cb5 (revert(views): calling order to initial view method (#7921))
 
     def get_serializer_context(self):
         context = super().get_serializer_context()
@@ -109,6 +127,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,8 +137,23 @@ class BaseTenantViewset(BaseViewSet):
             raise NotAuthenticated("Tenant ID is not present in token")
 
         user_id = str(request.user.id)
+<<<<<<< HEAD
         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
+>>>>>>> 6e7a32cb5 (revert(views): calling order to initial view method (#7921))
 
 
 class BaseUserViewset(BaseViewSet):
@@ -127,9 +162,11 @@ class BaseUserViewset(BaseViewSet):
             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 +174,22 @@ class BaseUserViewset(BaseViewSet):
         if tenant_id is None:
             raise NotAuthenticated("Tenant ID is not present in token")
 
+<<<<<<< HEAD
         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
+>>>>>>> 6e7a32cb5 (revert(views): calling order to initial view method (#7921))