chore: unify set config

Author: jfagoagas

api/src/backend/api/base_views.py

@@ -1,7 +1,6 @@
 import uuid
 
-from db_utils import tenant_transaction
-from django.db import connection, transaction
+from django.db import transaction
 from rest_framework import permissions
 from rest_framework.exceptions import NotAuthenticated
 from rest_framework.filters import SearchFilter
@@ -10,6 +9,7 @@
 from rest_framework_json_api.views import ModelViewSet
 from rest_framework_simplejwt.authentication import JWTAuthentication
 
+from api.db_utils import tenant_transaction
 from api.filters import CustomDjangoFilterBackend
 
 
@@ -75,8 +75,7 @@ def initial(self, request, *args, **kwargs):
             except ValueError:
                 raise ValidationError("User ID must be a valid UUID")
 
-            with connection.cursor() as cursor:
-                cursor.execute("SELECT set_config('api.user_id', %s, TRUE);", [user_id])
+            with tenant_transaction(value=user_id, parameter="api.user_id"):
                 return super().initial(request, *args, **kwargs)
 
         # TODO: DRY this when we have time

api/src/backend/api/db_utils.py

@@ -25,7 +25,7 @@
 POSTGRES_TENANT_VAR = "api.tenant_id"
 POSTGRES_USER_VAR = "api.user_id"
 
-SET_API_TENANT_ID_QUERY = "SELECT set_config('api.tenant_id', %s::text, TRUE);"
+SET_CONFIG_QUERY = "SELECT set_config(%s, %s::text, TRUE);"
 
 
 @contextmanager
@@ -48,15 +48,15 @@ def psycopg_connection(database_alias: str):
 
 
 @contextmanager
-def tenant_transaction(tenant_id: str):
+def tenant_transaction(value: str, parameter: str = "api.tenant_id"):
     with transaction.atomic():
         with connection.cursor() as cursor:
             try:
-                # just in case the tenant_id is an UUID object
-                uuid.UUID(str(tenant_id))
+                # just in case the tenant_id|user_id is an UUID object
+                uuid.UUID(str(value))
             except ValueError:
-                raise ValidationError("Tenant ID must be a valid UUID")
-            cursor.execute(SET_API_TENANT_ID_QUERY, [tenant_id])
+                raise ValidationError("Must be a valid UUID")
+            cursor.execute(SET_CONFIG_QUERY, [parameter, value])
             yield cursor
 
 

api/src/backend/api/decorators.py

@@ -4,7 +4,7 @@
 from django.db import connection, transaction
 from rest_framework_json_api.serializers import ValidationError
 
-from api.db_utils import SET_API_TENANT_ID_QUERY
+from api.db_utils import SET_CONFIG_QUERY
 
 
 def set_tenant(func):
@@ -52,7 +52,7 @@ def wrapper(*args, **kwargs):
         except ValueError:
             raise ValidationError("Tenant ID must be a valid UUID")
         with connection.cursor() as cursor:
-            cursor.execute(SET_API_TENANT_ID_QUERY, [tenant_id])
+            cursor.execute(SET_CONFIG_QUERY, ["api.tenant_id", tenant_id])
 
         return func(*args, **kwargs)