feat(rls-transaction): add retry for read replica connections (#9064)

Author: vicferpoy

.env

@@ -35,6 +35,8 @@ POSTGRES_DB=prowler_db
 # POSTGRES_REPLICA_USER=prowler
 # POSTGRES_REPLICA_PASSWORD=postgres
 # POSTGRES_REPLICA_DB=prowler_db
+# POSTGRES_REPLICA_MAX_ATTEMPTS=3
+# POSTGRES_REPLICA_RETRY_BASE_DELAY=0.5
 
 # Celery-Prowler task settings
 TASK_RETRY_DELAY_SECONDS=0.1

api/CHANGELOG.md

@@ -13,12 +13,17 @@ All notable changes to the **Prowler API** are documented in this file.
 - Support muting findings based on simple rules with custom reason [(#9051)](https://github.com/prowler-cloud/prowler/pull/9051)
 - Support C5 compliance framework for the GCP provider [(#9097)](https://github.com/prowler-cloud/prowler/pull/9097)
 
+---
+
 ## [1.14.1] (Prowler 5.13.1)
 
 ### Fixed
 - `/api/v1/overviews/providers` collapses data by provider type so the UI receives a single aggregated record per cloud family even when multiple accounts exist [(#9053)](https://github.com/prowler-cloud/prowler/pull/9053)
+- Added retry logic to database transactions to handle Aurora read replica connection failures during scale-down events [(#9064)](https://github.com/prowler-cloud/prowler/pull/9064)
 - Security Hub integrations stop failing when they read relationships via the replica by allowing replica relations and saving updates through the primary [(#9080)](https://github.com/prowler-cloud/prowler/pull/9080)
 
+---
+
 ## [1.14.0] (Prowler 5.13.0)
 
 ### Added

api/src/backend/api/db_utils.py

@@ -1,18 +1,35 @@
 import re
 import secrets
+import time
 import uuid
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 
+from celery.utils.log import get_task_logger
+from config.env import env
 from django.conf import settings
 from django.contrib.auth.models import BaseUserManager
-from django.db import DEFAULT_DB_ALIAS, connection, connections, models, transaction
+from django.db import (
+    DEFAULT_DB_ALIAS,
+    OperationalError,
+    connection,
+    connections,
+    models,
+    transaction,
+)
 from django_celery_beat.models import PeriodicTask
 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
 
-from api.db_router import get_read_db_alias, reset_read_db_alias, set_read_db_alias
+from api.db_router import (
+    READ_REPLICA_ALIAS,
+    get_read_db_alias,
+    reset_read_db_alias,
+    set_read_db_alias,
+)
+
+logger = get_task_logger(__name__)
 
 DB_USER = settings.DATABASES["default"]["USER"] if not settings.TESTING else "test"
 DB_PASSWORD = (
@@ -28,6 +45,9 @@
 POSTGRES_TENANT_VAR = "api.tenant_id"
 POSTGRES_USER_VAR = "api.user_id"
 
+REPLICA_MAX_ATTEMPTS = env.int("POSTGRES_REPLICA_MAX_ATTEMPTS", default=3)
+REPLICA_RETRY_BASE_DELAY = env.float("POSTGRES_REPLICA_RETRY_BASE_DELAY", default=0.5)
+
 SET_CONFIG_QUERY = "SELECT set_config(%s, %s::text, TRUE);"
 
 
@@ -71,24 +91,51 @@ def rls_transaction(
     if db_alias not in connections:
         db_alias = DEFAULT_DB_ALIAS
 
-    router_token = None
-    try:
-        if db_alias != DEFAULT_DB_ALIAS:
-            router_token = set_read_db_alias(db_alias)
-
-        with transaction.atomic(using=db_alias):
-            conn = connections[db_alias]
-            with conn.cursor() as cursor:
-                try:
-                    # just in case the value is a UUID object
-                    uuid.UUID(str(value))
-                except ValueError:
-                    raise ValidationError("Must be a valid UUID")
-                cursor.execute(SET_CONFIG_QUERY, [parameter, value])
-                yield cursor
-    finally:
-        if router_token is not None:
-            reset_read_db_alias(router_token)
+    alias = db_alias
+    is_replica = READ_REPLICA_ALIAS and alias == READ_REPLICA_ALIAS
+    max_attempts = REPLICA_MAX_ATTEMPTS if is_replica else 1
+
+    for attempt in range(1, max_attempts + 1):
+        router_token = None
+
+        # On final attempt, fallback to primary
+        if attempt == max_attempts and is_replica:
+            logger.warning(
+                f"RLS transaction failed after {attempt - 1} attempts on replica, "
+                f"falling back to primary DB"
+            )
+            alias = DEFAULT_DB_ALIAS
+
+        conn = connections[alias]
+        try:
+            if alias != DEFAULT_DB_ALIAS:
+                router_token = set_read_db_alias(alias)
+
+            with transaction.atomic(using=alias):
+                with conn.cursor() as cursor:
+                    try:
+                        # just in case the value is a UUID object
+                        uuid.UUID(str(value))
+                    except ValueError:
+                        raise ValidationError("Must be a valid UUID")
+                    cursor.execute(SET_CONFIG_QUERY, [parameter, value])
+                    yield cursor
+            return
+        except OperationalError as e:
+            # If on primary or max attempts reached, raise
+            if not is_replica or attempt == max_attempts:
+                raise
+
+            # Retry with exponential backoff
+            delay = REPLICA_RETRY_BASE_DELAY * (2 ** (attempt - 1))
+            logger.info(
+                f"RLS transaction failed on replica (attempt {attempt}/{max_attempts}), "
+                f"retrying in {delay}s. Error: {e}"
+            )
+            time.sleep(delay)
+        finally:
+            if router_token is not None:
+                reset_read_db_alias(router_token)
 
 
 class CustomUserManager(BaseUserManager):

api/src/backend/api/tests/integration/test_rls_transaction.py

@@ -0,0 +1,39 @@
+"""Tests for rls_transaction retry and fallback logic."""
+
+import pytest
+from django.db import DEFAULT_DB_ALIAS
+from rest_framework_json_api.serializers import ValidationError
+
+from api.db_utils import rls_transaction
+
+
+@pytest.mark.django_db
+class TestRLSTransaction:
+    """Simple integration tests for rls_transaction using real DB."""
+
+    @pytest.fixture
+    def tenant(self, tenants_fixture):
+        return tenants_fixture[0]
+
+    def test_success_on_primary(self, tenant):
+        """Basic: transaction succeeds on primary database."""
+        with rls_transaction(str(tenant.id), using=DEFAULT_DB_ALIAS) as cursor:
+            cursor.execute("SELECT 1")
+            result = cursor.fetchone()
+            assert result == (1,)
+
+    def test_invalid_uuid_raises_validation_error(self):
+        """Invalid UUID raises ValidationError before DB operations."""
+        with pytest.raises(ValidationError, match="Must be a valid UUID"):
+            with rls_transaction("not-a-uuid", using=DEFAULT_DB_ALIAS):
+                pass
+
+    def test_custom_parameter_name(self, tenant):
+        """Test custom RLS parameter name."""
+        custom_param = "api.custom_id"
+        with rls_transaction(
+            str(tenant.id), parameter=custom_param, using=DEFAULT_DB_ALIAS
+        ) as cursor:
+            cursor.execute("SELECT current_setting(%s, true)", [custom_param])
+            result = cursor.fetchone()
+            assert result == (str(tenant.id),)