feat(db): Update Django DB manager to use psycopg3 and connection poo…

…ling (#6541)

.env

@@ -92,3 +92,7 @@ jQIDAQAB
 # openssl rand -base64 32
 DJANGO_SECRETS_ENCRYPTION_KEY="oE/ltOhp/n1TdbHjVmzcjDPLcLA41CVI/4Rk+UB5ESc="
 DJANGO_BROKER_VISIBILITY_TIMEOUT=86400
+DJANGO_DB_CONNECTION_POOL_MIN_SIZE=4
+DJANGO_DB_CONNECTION_POOL_MAX_SIZE=10
+DJANGO_DB_CONNECTION_POOL_MAX_IDLE=36000
+DJANGO_DB_CONNECTION_POOL_MAX_LIFETIME=86400

api/.env.example

@@ -23,6 +23,10 @@ DJANGO_SECRETS_ENCRYPTION_KEY=""
 DJANGO_MANAGE_DB_PARTITIONS=[True|False]
 DJANGO_CELERY_DEADLOCK_ATTEMPTS=5
 DJANGO_BROKER_VISIBILITY_TIMEOUT=86400
+DJANGO_DB_CONNECTION_POOL_MIN_SIZE=4
+DJANGO_DB_CONNECTION_POOL_MAX_SIZE=10
+DJANGO_DB_CONNECTION_POOL_MAX_IDLE=36000
+DJANGO_DB_CONNECTION_POOL_MAX_LIFETIME=86400
 
 # PostgreSQL settings
 # If running django and celery on host, use 'localhost', else use 'postgres-db'

api/pyproject.toml

@@ -28,7 +28,7 @@ drf-spectacular = "0.27.2"
 drf-spectacular-jsonapi = "0.5.1"
 gunicorn = "23.0.0"
 prowler = "^5.0"
-psycopg2-binary = "2.9.9"
+psycopg = {extras = ["pool", "binary"], version = "3.2.3"}
 pytest-celery = {extras = ["redis"], version = "^1.0.1"}
 # Needed for prowler compatibility
 python = ">=3.11,<3.13"

api/src/backend/api/db_utils.py

@@ -6,8 +6,10 @@
 from django.conf import settings
 from django.contrib.auth.models import BaseUserManager
 from django.db import connection, models, transaction
-from psycopg2 import connect as psycopg2_connect
-from psycopg2.extensions import AsIs, new_type, register_adapter, register_type
+from psycopg import connect as psycopg_connect
+from psycopg.adapt import Dumper
+from psycopg.types import TypeInfo
+from psycopg.types.string import TextLoader
 from rest_framework_json_api.serializers import ValidationError
 
 DB_USER = settings.DATABASES["default"]["USER"] if not settings.TESTING else "test"
@@ -20,6 +22,7 @@
 DB_PROWLER_PASSWORD = (
     settings.DATABASES["prowler_user"]["PASSWORD"] if not settings.TESTING else "test"
 )
+
 TASK_RUNNER_DB_TABLE = "django_celery_results_taskresult"
 POSTGRES_TENANT_VAR = "api.tenant_id"
 POSTGRES_USER_VAR = "api.user_id"
@@ -29,21 +32,25 @@
 
 @contextmanager
 def psycopg_connection(database_alias: str):
-    psycopg2_connection = None
+    """
+    Context manager returning a psycopg 3 connection
+    for the specified 'database_alias' in Django settings.
+    """
+    pg_conn = None
     try:
         admin_db = settings.DATABASES[database_alias]
 
-        psycopg2_connection = psycopg2_connect(
+        pg_conn = psycopg_connect(
             dbname=admin_db["NAME"],
             user=admin_db["USER"],
             password=admin_db["PASSWORD"],
             host=admin_db["HOST"],
             port=admin_db["PORT"],
         )
-        yield psycopg2_connection
+        yield pg_conn
     finally:
-        if psycopg2_connection is not None:
-            psycopg2_connection.close()
+        if pg_conn is not None:
+            pg_conn.close()
 
 
 @contextmanager
@@ -59,7 +66,7 @@ def rls_transaction(value: str, parameter: str = POSTGRES_TENANT_VAR):
     with transaction.atomic():
         with connection.cursor() as cursor:
             try:
-                # just in case the value is an UUID object
+                # Just in case the value is a UUID object
                 uuid.UUID(str(value))
             except ValueError:
                 raise ValidationError("Must be a valid UUID")
@@ -187,32 +194,24 @@ def __str__(self):
         return self.value
 
 
-def enum_adapter(enum_obj):
-    return AsIs(f"'{enum_obj.value}'::{enum_obj.__class__.enum_type_name}")
-
-
-def get_enum_oid(connection, enum_type_name: str):
-    with connection.cursor() as cursor:
-        cursor.execute("SELECT oid FROM pg_type WHERE typname = %s;", (enum_type_name,))
-        result = cursor.fetchone()
-    if result is None:
-        raise ValueError(f"Enum type '{enum_type_name}' not found")
-    return result[0]
-
+def register_enum(apps, schema_editor, enum_class):
+    """
+    psycopg 3 approach: register a loader + dumper for the given enum_class,
+    so we can read/write the custom Postgres ENUM seamlessly.
+    """
+    with psycopg_connection(schema_editor.connection.alias) as conn:
+        ti = TypeInfo.fetch(conn, enum_class.enum_type_name)
 
-def register_enum(apps, schema_editor, enum_class):  # noqa: F841
-    with psycopg_connection(schema_editor.connection.alias) as connection:
-        enum_oid = get_enum_oid(connection, enum_class.enum_type_name)
-        enum_instance = new_type(
-            (enum_oid,),
-            enum_class.enum_type_name,
-            lambda value, cur: value,  # noqa: F841
-        )
-        register_type(enum_instance, connection)
-        register_adapter(enum_class, enum_adapter)
+        class EnumLoader(TextLoader):
+            def load(self, data):
+                return data
 
+        class EnumDumper(Dumper):
+            def dump(self, obj):
+                return f"'{obj.value}'::{obj.__class__.enum_type_name}"
 
-# Postgres enum definition for member role
+        conn.adapters.register_loader(ti.oid, EnumLoader)
+        conn.adapters.register_dumper(enum_class, EnumDumper)
 
 
 class MemberRoleEnum(EnumType):

api/src/backend/config/django/base.py

@@ -115,6 +115,12 @@
 
 DATABASE_ROUTERS = ["api.db_router.MainRouter"]
 
+# Database connection pool
+DB_CP_MIN_SIZE = env.int("DJANGO_DB_CONNECTION_POOL_MIN_SIZE", 4)
+DB_CP_MAX_SIZE = env.int("DJANGO_DB_CONNECTION_POOL_MAX_SIZE", 10)
+DB_CP_MAX_IDLE = env.int("DJANGO_DB_CONNECTION_POOL_MAX_IDLE", 36000)
+DB_CP_MAX_LIFETIME = env.int("DJANGO_DB_CONNECTION_POOL_MAX_LIFETIME", 86400)
+
 
 # Password validation
 # https://docs.djangoproject.com/en/5.0/ref/settings/#auth-password-validators

api/src/backend/config/django/devel.py

@@ -1,7 +1,6 @@
 from config.django.base import *  # noqa
 from config.env import env
 
-
 DEBUG = env.bool("DJANGO_DEBUG", default=True)
 ALLOWED_HOSTS = env.list("DJANGO_ALLOWED_HOSTS", default=["*"])
 
@@ -14,6 +13,14 @@
         "PASSWORD": env("POSTGRES_PASSWORD", default="prowler"),
         "HOST": env("POSTGRES_HOST", default="postgres-db"),
         "PORT": env("POSTGRES_PORT", default="5432"),
+        "OPTIONS": {
+            "pool": {
+                "min_size": DB_CP_MIN_SIZE,  # noqa: F405
+                "max_size": DB_CP_MAX_SIZE,  # noqa: F405
+                "max_idle": DB_CP_MAX_IDLE,  # noqa: F405
+                "max_lifetime": DB_CP_MAX_LIFETIME,  # noqa: F405
+            }
+        },
     },
     "admin": {
         "ENGINE": "psqlextra.backend",
@@ -22,6 +29,14 @@
         "PASSWORD": env("POSTGRES_ADMIN_PASSWORD", default="S3cret"),
         "HOST": env("POSTGRES_HOST", default="postgres-db"),
         "PORT": env("POSTGRES_PORT", default="5432"),
+        "OPTIONS": {
+            "pool": {
+                "min_size": DB_CP_MIN_SIZE,  # noqa: F405
+                "max_size": DB_CP_MAX_SIZE,  # noqa: F405
+                "max_idle": DB_CP_MAX_IDLE,  # noqa: F405
+                "max_lifetime": DB_CP_MAX_LIFETIME,  # noqa: F405
+            }
+        },
     },
 }
 DATABASES["default"] = DATABASES["prowler_user"]

api/src/backend/config/django/production.py

@@ -1,7 +1,6 @@
 from config.django.base import *  # noqa
 from config.env import env
 
-
 DEBUG = env.bool("DJANGO_DEBUG", default=False)
 ALLOWED_HOSTS = env.list("DJANGO_ALLOWED_HOSTS", default=["localhost", "127.0.0.1"])
 
@@ -15,6 +14,14 @@
         "PASSWORD": env("POSTGRES_PASSWORD"),
         "HOST": env("POSTGRES_HOST"),
         "PORT": env("POSTGRES_PORT"),
+        "OPTIONS": {
+            "pool": {
+                "min_size": DB_CP_MIN_SIZE,  # noqa: F405
+                "max_size": DB_CP_MAX_SIZE,  # noqa: F405
+                "max_idle": DB_CP_MAX_IDLE,  # noqa: F405
+                "max_lifetime": DB_CP_MAX_LIFETIME,  # noqa: F405
+            }
+        },
     },
     "admin": {
         "ENGINE": "psqlextra.backend",
@@ -23,6 +30,14 @@
         "PASSWORD": env("POSTGRES_ADMIN_PASSWORD"),
         "HOST": env("POSTGRES_HOST"),
         "PORT": env("POSTGRES_PORT"),
+        "OPTIONS": {
+            "pool": {
+                "min_size": DB_CP_MIN_SIZE,  # noqa: F405
+                "max_size": DB_CP_MAX_SIZE,  # noqa: F405
+                "max_idle": DB_CP_MAX_IDLE,  # noqa: F405
+                "max_lifetime": DB_CP_MAX_LIFETIME,  # noqa: F405
+            }
+        },
     },
 }
 DATABASES["default"] = DATABASES["prowler_user"]

api/src/backend/config/django/testing.py

@@ -1,7 +1,6 @@
 from config.django.base import *  # noqa
 from config.env import env
 
-
 DEBUG = env.bool("DJANGO_DEBUG", default=False)
 ALLOWED_HOSTS = env.list("DJANGO_ALLOWED_HOSTS", default=["localhost", "127.0.0.1"])
 
@@ -14,6 +13,14 @@
         "PASSWORD": env("POSTGRES_PASSWORD", default="postgres"),
         "HOST": env("POSTGRES_HOST", default="localhost"),
         "PORT": env("POSTGRES_PORT", default="5432"),
+        "OPTIONS": {
+            "pool": {
+                "min_size": DB_CP_MIN_SIZE,  # noqa: F405
+                "max_size": DB_CP_MAX_SIZE,  # noqa: F405
+                "max_idle": DB_CP_MAX_IDLE,  # noqa: F405
+                "max_lifetime": DB_CP_MAX_LIFETIME,  # noqa: F405
+            }
+        },
     },
 }