ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-10-09 21:14:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(db): Update Django DB manager to use psycopg3 and connection pooling (#6541)

Browse Source
This commit is contained in:
Víctor Fernández Poyatos
2025-01-16 15:29:02 +01:00
committed by GitHub
parent 98d9256f92
commit 8821a91f3f
10 changed files with 195 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/backend/api/db_utils.py
+30 -31
View File
@@ -6,8 +6,10 @@ from datetime import datetime, timedelta, timezone
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_USER = (
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 @@ SET_CONFIG_QUERY = "SELECT set_config(%s, %s::text, TRUE);"
@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 @@ class EnumType:
return self.value
def enum_adapter(enum_obj):
return AsIs(f"'{enum_obj.value}'::{enum_obj.__class__.enum_type_name}")
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)
class EnumLoader(TextLoader):
def load(self, data):
return data
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]
class EnumDumper(Dumper):
def dump(self, obj):
return f"'{obj.value}'::{obj.__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)
# 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
+6
View File
@@ -115,6 +115,12 @@ DJANGO_GUID = {
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
+16 -1
View File
@@ -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 @@ DATABASES = {
"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 @@ DATABASES = {
"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
+16 -1
View File
@@ -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 @@ DATABASES = {
"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 @@ DATABASES = {
"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
+8 -1
View File
@@ -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 @@ DATABASES = {
"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
}
},
},
}