Merge branch 'master' into PRWLR-6072-scheduled-scans-must-appear-in-…

…scans-right-after-creating-a-provider

.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/docker-entrypoint.sh

@@ -28,7 +28,7 @@ start_prod_server() {
 
 start_worker() {
   echo "Starting the worker..."
-  poetry run python -m celery -A config.celery worker -l "${DJANGO_LOGGING_LEVEL:-info}" -Q celery,scans -E
+  poetry run python -m celery -A config.celery worker -l "${DJANGO_LOGGING_LEVEL:-info}" -Q celery,scans -E --max-tasks-per-child 1
 }
 
 start_worker_beat() {

api/pyproject.toml

@@ -8,7 +8,7 @@ description = "Prowler's API (Django/DRF)"
 license = "Apache-2.0"
 name = "prowler-api"
 package-mode = false
-version = "1.3.0"
+version = "1.4.0"
 
 [tool.poetry.dependencies]
 celery = {extras = ["pytest"], version = "^5.4.0"}
@@ -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"
@@ -37,6 +37,7 @@ uuid6 = "2024.7.10"
 [tool.poetry.group.dev.dependencies]
 bandit = "1.7.9"
 coverage = "7.5.4"
+django-silk = "5.3.2"
 docker = "7.1.0"
 freezegun = "1.5.1"
 mypy = "1.10.1"

api/src/backend/api/db_router.py

@@ -4,13 +4,17 @@ class MainRouter:
 
     def db_for_read(self, model, **hints):  # noqa: F841
         model_table_name = model._meta.db_table
-        if model_table_name.startswith("django_"):
+        if model_table_name.startswith("django_") or model_table_name.startswith(
+            "silk_"
+        ):
             return self.admin_db
         return None
 
     def db_for_write(self, model, **hints):  # noqa: F841
         model_table_name = model._meta.db_table
-        if model_table_name.startswith("django_"):
+        if model_table_name.startswith("django_") or model_table_name.startswith(
+            "silk_"
+        ):
             return self.admin_db
         return None
 

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/api/filters.py

@@ -319,26 +319,27 @@ class FindingFilter(FilterSet):
         field_name="resources__type", lookup_expr="icontains"
     )
 
-    resource_tag_key = CharFilter(field_name="resources__tags__key")
-    resource_tag_key__in = CharInFilter(
-        field_name="resources__tags__key", lookup_expr="in"
-    )
-    resource_tag_key__icontains = CharFilter(
-        field_name="resources__tags__key", lookup_expr="icontains"
-    )
-    resource_tag_value = CharFilter(field_name="resources__tags__value")
-    resource_tag_value__in = CharInFilter(
-        field_name="resources__tags__value", lookup_expr="in"
-    )
-    resource_tag_value__icontains = CharFilter(
-        field_name="resources__tags__value", lookup_expr="icontains"
-    )
-    resource_tags = CharInFilter(
-        method="filter_resource_tag",
-        lookup_expr="in",
-        help_text="Filter by resource tags `key:value` pairs.\nMultiple values may be "
-        "separated by commas.",
-    )
+    # Temporarily disabled until we implement tag filtering in the UI
+    # resource_tag_key = CharFilter(field_name="resources__tags__key")
+    # resource_tag_key__in = CharInFilter(
+    #     field_name="resources__tags__key", lookup_expr="in"
+    # )
+    # resource_tag_key__icontains = CharFilter(
+    #     field_name="resources__tags__key", lookup_expr="icontains"
+    # )
+    # resource_tag_value = CharFilter(field_name="resources__tags__value")
+    # resource_tag_value__in = CharInFilter(
+    #     field_name="resources__tags__value", lookup_expr="in"
+    # )
+    # resource_tag_value__icontains = CharFilter(
+    #     field_name="resources__tags__value", lookup_expr="icontains"
+    # )
+    # resource_tags = CharInFilter(
+    #     method="filter_resource_tag",
+    #     lookup_expr="in",
+    #     help_text="Filter by resource tags `key:value` pairs.\nMultiple values may be "
+    #     "separated by commas.",
+    # )
 
     scan = UUIDFilter(method="filter_scan_id")
     scan__in = UUIDInFilter(method="filter_scan_id_in")
@@ -374,12 +375,6 @@ class Meta:
             },
         }
 
-    @property
-    def qs(self):
-        # Force distinct results to prevent duplicates with many-to-many relationships
-        parent_qs = super().qs
-        return parent_qs.distinct()
-
     #  Convert filter values to UUIDv7 values for use with partitioning
     def filter_scan_id(self, queryset, name, value):
         try:

api/src/backend/api/migrations/0007_scan_and_scan_summaries_indexes.py

@@ -0,0 +1,25 @@
+# Generated by Django 5.1.5 on 2025-01-28 15:03
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("api", "0006_findings_first_seen"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="scan",
+            index=models.Index(
+                fields=["tenant_id", "provider_id", "state", "inserted_at"],
+                name="scans_prov_state_insert_idx",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="scansummary",
+            index=models.Index(
+                fields=["tenant_id", "scan_id"], name="scan_summaries_tenant_scan_idx"
+            ),
+        ),
+    ]

api/src/backend/api/models.py

@@ -432,6 +432,10 @@ class Meta(RowLevelSecurityProtectedModel.Meta):
                 fields=["provider", "state", "trigger", "scheduled_at"],
                 name="scans_prov_state_trig_sche_idx",
             ),
+            models.Index(
+                fields=["tenant_id", "provider_id", "state", "inserted_at"],
+                name="scans_prov_state_insert_idx",
+            ),
         ]
 
     class JSONAPIMeta:
@@ -1104,6 +1108,12 @@ class Meta(RowLevelSecurityProtectedModel.Meta):
                 statements=["SELECT", "INSERT", "UPDATE", "DELETE"],
             ),
         ]
+        indexes = [
+            models.Index(
+                fields=["tenant_id", "scan_id"],
+                name="scan_summaries_tenant_scan_idx",
+            )
+        ]
 
     class JSONAPIMeta:
         resource_name = "scan-summaries"