From 8a2cfea6776cd03e64f4c3f3d5ce12f293798738 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Fern=C3=A1ndez=20Poyatos?= Date: Wed, 21 Aug 2024 10:10:01 +0200 Subject: [PATCH] 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 | 19 +++++++++++++- src/backend/api/middleware.py | 33 ------------------------ src/backend/api/tests/test_middleware.py | 30 +-------------------- src/backend/config/django/base.py | 1 - 4 files changed, 19 insertions(+), 64 deletions(-) diff --git a/src/backend/api/base_views.py b/src/backend/api/base_views.py index f353d1a688..5c30218196 100644 --- a/src/backend/api/base_views.py +++ b/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") diff --git a/src/backend/api/middleware.py b/src/backend/api/middleware.py index 7554d5823d..5d4f02537d 100644 --- a/src/backend/api/middleware.py +++ b/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]) diff --git a/src/backend/api/tests/test_middleware.py b/src/backend/api/tests/test_middleware.py index 88c94beb00..44e91bf200 100644 --- a/src/backend/api/tests/test_middleware.py +++ b/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) diff --git a/src/backend/config/django/base.py b/src/backend/config/django/base.py index bb7e3efaef..b8f2d66797 100644 --- a/src/backend/config/django/base.py +++ b/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", ]