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

revert: Update Django DB manager to use psycopg3 and connection pooling (#6719)

Browse Source 
Co-authored-by: Víctor Fernández Poyatos <victor@prowler.com>
This commit is contained in:
Prowler Bot
2025-01-28 22:33:32 +05:45
committed by GitHub
co-authored by Víctor Fernández Poyatos
parent 1fe125867c
commit 9630f2242a
10 changed files with 444 additions and 532 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/.env.example
-4
View File
@@ -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/poetry.lock
Generated
+411 -441
View File
File diff suppressed because it is too large Load Diff
api/pyproject.toml
+1 -1
View File
@@ -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 = "v5.2"}
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
+31 -30
View File
@@ -6,10 +6,8 @@ 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 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_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"
@@ -32,25 +29,21 @@ SET_CONFIG_QUERY = "SELECT set_config(%s, %s::text, TRUE);"
@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 @@ class EnumType:
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]
conn.adapters.register_loader(ti.oid, EnumLoader)
conn.adapters.register_dumper(enum_class, EnumDumper)
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
class MemberRoleEnum(EnumType):
api/src/backend/config/django/base.py
-6
View File
@@ -115,12 +115,6 @@ 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
View File
@@ -13,14 +13,6 @@ 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",
@@ -29,14 +21,6 @@ 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
View File
@@ -14,14 +14,6 @@ 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",
@@ -30,14 +22,6 @@ 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
View File
@@ -13,14 +13,6 @@ 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
}
},
},
}