feat/PRWLR-4368 Remove TenantMiddleware and set tenant_id at view level (#31)

* feat(API): PRWLR-4368 remove TenantMiddleware in favour of transaction based setup

* feat(API): PRWLR-4368 override initial request method to perform atomic transactions on RLS viewsets

src/backend/api/base_views.py

@@ -1,6 +1,10 @@
+import uuid
+
+from django.db import transaction, connection
 from rest_framework.exceptions import NotAuthenticated
 from rest_framework.filters import SearchFilter
 from rest_framework_json_api import filters
+from rest_framework_json_api.serializers import ValidationError
 from rest_framework_json_api.views import ModelViewSet
 
 from api.filters import CustomDjangoFilterBackend
@@ -26,12 +30,25 @@ class BaseViewSet(ModelViewSet):
 
 class BaseRLSViewSet(BaseViewSet):
     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 "X-Tenant-ID" not in request.headers:
             # This will return a 403 until we implement authentication/authorization
             # https://www.django-rest-framework.org/api-guide/authentication/#unauthorized-and-forbidden-responses
             raise NotAuthenticated("X-Tenant-ID header is required")
 
+        tenant_id = request.headers["X-Tenant-ID"]
+
+        try:
+            uuid.UUID(tenant_id)
+        except ValueError:
+            raise ValidationError("X-Tenant-ID header must be a valid UUID")
+
+        with transaction.atomic():
+            with connection.cursor() as cursor:
+                cursor.execute("SET api.tenant_id = %s", [tenant_id])
+            return super().initial(request, *args, **kwargs)
+
     def get_serializer_context(self):
         context = super().get_serializer_context()
         tenant_id = self.request.headers.get("X-Tenant-ID")

src/backend/api/middleware.py

@@ -1,7 +1,6 @@
 import logging
 import time
 
-from django.db import connection
 from django.http import HttpRequest
 
 from config.custom_logging import BackendLogger
@@ -52,35 +51,3 @@ class APILoggingMiddleware:
         )
 
         return response
-
-
-class TenantMiddleware:
-    """
-    Middleware to handle setting the tenant ID for row-level security (RLS) in the current session.
-
-    This middleware extracts the tenant ID from the request headers and sets it in the session for RLS.
-    """
-
-    def __init__(self, get_response):
-        self.get_response = get_response
-
-    def __call__(self, request):
-        tenant_id = extract_tenant_id(request)
-
-        # TODO Define exception handling for POST/PATCH requests without tenant ID
-        if tenant_id:
-            self.set_tenant_id_in_session(tenant_id)
-
-        response = self.get_response(request)
-        return response
-
-    @staticmethod
-    def set_tenant_id_in_session(tenant_id: str):
-        """
-        Set the tenant ID in the session for RLS using a raw SQL query.
-
-        Args:
-            tenant_id (str): The tenant ID to set in the session. It is a UUID string.
-        """
-        with connection.cursor() as cursor:
-            cursor.execute("SET api.tenant_id = %s", [tenant_id])

src/backend/api/tests/test_middleware.py

@@ -1,9 +1,8 @@
 from unittest.mock import patch, Mock, ANY
 
 import pytest
-from django.http import HttpResponse
 
-from api.middleware import extract_tenant_id, TenantMiddleware
+from api.middleware import extract_tenant_id
 
 
 @pytest.mark.django_db
@@ -36,33 +35,6 @@ def test_api_logger_middleware(mock_get_logger, client):
     assert isinstance(mock_logger.info.call_args[1]["extra"]["duration"], float)
 
 
-@patch("api.middleware.TenantMiddleware.set_tenant_id_in_session")
-def test_tenant_middleware(mock_set_tenant, client):
-    tenant_id = "12646005-9067-4d2a-a098-8bb378604362"
-    tenant_header = {"X-Tenant-ID": tenant_id}
-
-    response = client.get("/testing", headers=tenant_header)
-
-    mock_set_tenant.assert_called_once_with(tenant_id)
-    assert isinstance(response, HttpResponse)
-
-
-@patch("api.middleware.connection.cursor")
-def test_tenant_middleware_set_tenant_id_in_session(cursor_mock):
-    cursor_mock.return_value.__enter__.return_value = cursor_mock
-    cursor_mock.execute.return_value = None
-
-    tenant_id_postgres_variable = "api.tenant_id"
-    tenant_id = "12646005-9067-4d2a-a098-8bb378604362"
-    tenant_middleware = TenantMiddleware(Mock())
-
-    tenant_middleware.set_tenant_id_in_session(tenant_id)
-
-    cursor_mock.execute.assert_called_once_with(
-        f"SET {tenant_id_postgres_variable} = %s", [tenant_id]
-    )
-
-
 def test_extract_tenant_id():
     mock_request = Mock()
     extract_tenant_id(mock_request)

src/backend/config/django/base.py

@@ -34,7 +34,6 @@ MIDDLEWARE = [
     "django.contrib.auth.middleware.AuthenticationMiddleware",
     "django.contrib.messages.middleware.MessageMiddleware",
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
-    "api.middleware.TenantMiddleware",
     "api.middleware.APILoggingMiddleware",
 ]