revert: Update Django DB manager to use psycopg3 and connection pooli…

…ng (#6717)

.env

@@ -92,7 +92,3 @@ 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,10 +23,6 @@ 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 = {git = "https://github.com/prowler-cloud/prowler.git", branch = "master"}
-psycopg = {extras = ["pool", "binary"], version = "3.2.3"}
+psycopg2-binary = "2.9.9"
 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,10 +6,8 @@
 from django.conf import settings
 from django.contrib.auth.models import BaseUserManager
 from django.db import connection, models, transaction
-from psycopg import connect as psycopg_connect
-from psycopg.adapt import Dumper
-from psycopg.types import TypeInfo
-from psycopg.types.string import TextLoader
+from psycopg2 import connect as psycopg2_connect
+from psycopg2.extensions import AsIs, new_type, register_adapter, register_type
 from rest_framework_json_api.serializers import ValidationError
 
 DB_USER = settings.DATABASES["default"]["USER"] if not settings.TESTING else "test"
@@ -22,7 +20,6 @@
 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"
@@ -32,25 +29,21 @@
 
 @contextmanager
 def psycopg_connection(database_alias: str):
-    """
-    Context manager returning a psycopg 3 connection
-    for the specified 'database_alias' in Django settings.
-    """
-    pg_conn = None
+    psycopg2_connection = None
     try:
         admin_db = settings.DATABASES[database_alias]
 
-        pg_conn = psycopg_connect(
+        psycopg2_connection = psycopg2_connect(
             dbname=admin_db["NAME"],
             user=admin_db["USER"],
             password=admin_db["PASSWORD"],
             host=admin_db["HOST"],
             port=admin_db["PORT"],
         )
-        yield pg_conn
+        yield psycopg2_connection
     finally:
-        if pg_conn is not None:
-            pg_conn.close()
+        if psycopg2_connection is not None:
+            psycopg2_connection.close()
 
 
 @contextmanager
@@ -66,7 +59,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 a UUID object
+                # just in case the value is an UUID object
                 uuid.UUID(str(value))
             except ValueError:
                 raise ValidationError("Must be a valid UUID")
@@ -194,24 +187,32 @@ def __str__(self):
         return self.value
 
 
-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 enum_adapter(enum_obj):
+    return AsIs(f"'{enum_obj.value}'::{enum_obj.__class__.enum_type_name}")
 
-        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}"
+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):  # 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)
+
 
-        conn.adapters.register_loader(ti.oid, EnumLoader)
-        conn.adapters.register_dumper(enum_class, EnumDumper)
+# Postgres enum definition for member role
 
 
 class MemberRoleEnum(EnumType):

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

@@ -115,12 +115,6 @@
 
 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

@@ -13,14 +13,6 @@
         "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",
@@ -29,14 +21,6 @@
         "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

@@ -14,14 +14,6 @@
         "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",
@@ -30,14 +22,6 @@
         "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

@@ -13,14 +13,6 @@
         "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
-            }
-        },
     },
 }