ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-23 04:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): add query-level retry with primary fallback to rls_transaction via execute_wrapper (#10379)

Browse Source 
Co-authored-by: Pepe Fagoaga <pepe@prowler.com>
This commit is contained in:
Josema Camacho
2026-07-09 17:47:03 +02:00
committed by GitHub
parent 6f72785a28
commit b44f8ebf5f
5 changed files with 905 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/changelog.d/rls-transaction-primary-fallback.fixed.md
+1
View File
@@ -0,0 +1 @@
`rls_transaction` now falls back directly to the primary DB for connection-level mid-query read replica failures via `execute_wrapper`, reducing non-streaming read crashes during replica recovery
api/src/backend/api/db_utils.py
+273 -50
View File
@@ -2,7 +2,7 @@ import re
import secrets
import time
import uuid
from contextlib import contextmanager
from contextlib import ExitStack, contextmanager, nullcontext
from datetime import UTC, datetime, timedelta
from api.db_router import (
@@ -48,6 +48,140 @@ 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);"
SET_TRANSACTION_READ_ONLY_QUERY = "SET TRANSACTION READ ONLY;"
REPLICA_CONNECTION_SQLSTATE_PREFIXES = ("08",)
REPLICA_CONNECTION_SQLSTATES = {"57P01", "57P02", "57P03"}
REPLICA_NON_FAILOVER_SQLSTATES = {"57014", "40001", "40P01"}
REPLICA_CONNECTION_ERROR_MESSAGES = (
"ssl syscall",
"eof detected",
"server closed the connection",
"connection already closed",
"connection not open",
"could not connect to server",
"connection refused",
"connection reset",
"connection timed out",
"lost synchronization",
"terminating connection",
"database system is starting up",
"database system is shutting down",
"database system is in recovery mode",
)
REPLICA_NON_FAILOVER_ERROR_MESSAGES = (
"canceling statement due to user request",
"deadlock detected",
"could not serialize access",
)
def _iter_exception_chain(error: BaseException):
seen = set()
pending = [error]
while pending:
current = pending.pop(0)
if current is None or id(current) in seen:
continue
seen.add(id(current))
yield current
cause = getattr(current, "__cause__", None)
context = getattr(current, "__context__", None)
if cause is not None:
pending.append(cause)
if context is not None:
pending.append(context)
for arg in getattr(current, "args", ()):
if isinstance(arg, BaseException):
pending.append(arg)
def _get_exception_sqlstate(error: BaseException) -> str | None:
for attr in ("pgcode", "sqlstate"):
sqlstate = getattr(error, attr, None)
if sqlstate:
return sqlstate
diag = getattr(error, "diag", None)
if diag is not None:
sqlstate = getattr(diag, "sqlstate", None)
if sqlstate:
return sqlstate
return None
def _is_replica_connection_failure(error: BaseException) -> bool:
"""
Return True only for replica failures where retrying on primary is safe.
Query cancellations, serialization failures, and deadlocks should surface to
callers because replaying them can hide real query or concurrency problems.
"""
messages = []
sqlstates = set()
for chained_error in _iter_exception_chain(error):
sqlstate = _get_exception_sqlstate(chained_error)
if sqlstate:
sqlstates.add(sqlstate)
messages.append(str(chained_error).lower())
if sqlstates & REPLICA_NON_FAILOVER_SQLSTATES:
return False
if any(
sqlstate.startswith(REPLICA_CONNECTION_SQLSTATE_PREFIXES)
or sqlstate in REPLICA_CONNECTION_SQLSTATES
for sqlstate in sqlstates
):
return True
message = " ".join(messages)
if any(marker in message for marker in REPLICA_NON_FAILOVER_ERROR_MESSAGES):
return False
return any(marker in message for marker in REPLICA_CONNECTION_ERROR_MESSAGES)
def _strip_leading_sql_comments(sql: str) -> str:
if not isinstance(sql, str):
return ""
sql_text = sql.lstrip()
while True:
if sql_text.startswith("--"):
newline_index = sql_text.find("\n")
if newline_index == -1:
return ""
sql_text = sql_text[newline_index + 1 :].lstrip()
continue
if sql_text.startswith("/*"):
comment_end_index = sql_text.find("*/", 2)
if comment_end_index == -1:
return ""
sql_text = sql_text[comment_end_index + 2 :].lstrip()
continue
return sql_text
def _is_safe_primary_replay(sql: str, many: bool) -> bool:
if many:
return False
sql_text = _strip_leading_sql_comments(sql)
if not re.match(r"(?is)^SELECT\b", sql_text):
return False
return not any(
re.search(pattern, sql_text, re.IGNORECASE | re.DOTALL)
for pattern in (
r"\bINTO\b",
r"\bFOR\s+(?:NO\s+KEY\s+)?UPDATE\b",
r"\bFOR\s+(?:KEY\s+)?SHARE\b",
)
)
@contextmanager
@@ -77,14 +211,36 @@ def rls_transaction(
retry_on_replica: bool = True,
):
"""
Creates a new database transaction setting the given configuration value for Postgres RLS. It validates the
if the value is a valid UUID.
Context manager that opens an RLS-scoped database transaction.
Sets a Postgres configuration variable (``set_config``) so that Row-Level
Security policies can filter by tenant. When *using* points to a read
replica and *retry_on_replica* is True, replica failures are handled in two
places:
1. **Pre-yield** (connection-setup failures): the function retries
up to ``REPLICA_MAX_ATTEMPTS`` times on the replica, then falls
back to the primary DB.
2. **Post-yield** (mid-query failures): an ``execute_wrapper``
intercepts connection-level ``OperationalError`` during
``cursor.execute()`` calls and falls back directly to the primary DB
for single ``SELECT`` statements. The primary fallback transaction is
read-only, and unsafe statements keep raising the original error.
The wrapper swaps the inner cursor so ``fetchall()`` / ``fetchone()``
read from the new connection transparently.
Limitation: server-side cursors (``.iterator()``) fetch rows via
``fetchmany()``, which the wrapper does not intercept. Call sites
that iterate large result sets with ``.iterator()`` on the replica
should add their own retry logic.
Args:
value (str): Database configuration parameter value.
parameter (str): Database configuration parameter name, by default is 'api.tenant_id'.
using (str | None): Optional database alias to run the transaction against. Defaults to the
active read alias (if any) or Django's default connection.
value: Database configuration parameter value (must be a valid UUID).
parameter: Database configuration parameter name.
using: Optional database alias. Defaults to the active read
alias or Django's default connection.
retry_on_replica: Whether replica setup failures can retry and
connection-level mid-query failures can fall back to primary.
"""
requested_alias = using or get_read_db_alias()
db_alias = requested_alias or DEFAULT_DB_ALIAS
@@ -92,54 +248,121 @@ def rls_transaction(
db_alias = DEFAULT_DB_ALIAS
alias = db_alias
is_replica = READ_REPLICA_ALIAS and alias == READ_REPLICA_ALIAS
max_attempts = REPLICA_MAX_ATTEMPTS if is_replica and retry_on_replica else 1
is_replica = bool(READ_REPLICA_ALIAS and alias == READ_REPLICA_ALIAS)
can_failover = is_replica and retry_on_replica
replica_alias = alias # captured before the loop mutates alias
max_attempts = (REPLICA_MAX_ATTEMPTS + 1) if can_failover else 1
for attempt in range(1, max_attempts + 1):
router_token = None
yielded_cursor = False
# State shared between the generator and the _query_failover closure.
# The fallback transaction.atomic() is registered into fallback_stack
# via enter_context so its __exit__ runs when the outer with-ExitStack
# block exits, with the right exc_info. No manual __enter__/__exit__.
_fallback = {"succeeded": False, "token": None, "caller_exited_cleanly": False}
# 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
with ExitStack() as fallback_stack:
conn = connections[alias]
try:
if alias != DEFAULT_DB_ALIAS:
router_token = set_read_db_alias(alias)
def _query_failover(execute, sql, params, many, context):
"""execute_wrapper: replay failed replica queries on the primary DB."""
try:
return execute(sql, params, many, context)
except OperationalError as err:
if not _is_replica_connection_failure(err):
raise
if not _is_safe_primary_replay(sql, many):
raise
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])
yielded_cursor = True
yield cursor
return
except OperationalError as e:
if yielded_cursor:
raise
# If on primary or max attempts reached, raise
if not is_replica or attempt == max_attempts:
raise
try:
connections[replica_alias].close()
except Exception:
pass # Best-effort; connection may already be dead
# 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)
logger.warning(
"Mid-query replica connection failure, falling back to primary DB"
)
primary = connections[DEFAULT_DB_ALIAS]
primary.ensure_connection()
fallback_stack.enter_context(transaction.atomic(using=DEFAULT_DB_ALIAS))
fallback_cursor = primary.cursor()
fallback_stack.callback(fallback_cursor.close)
fallback_cursor.execute(SET_TRANSACTION_READ_ONLY_QUERY)
fallback_cursor.execute(SET_CONFIG_QUERY, [parameter, value])
_fallback["token"] = set_read_db_alias(DEFAULT_DB_ALIAS)
fallback_cursor.execute(sql, params)
context["cursor"].db = primary
context["cursor"].cursor = fallback_cursor.cursor
_fallback["succeeded"] = True
return None
for attempt in range(1, max_attempts + 1):
router_token = None
yielded_cursor = False
# On final attempt, fall back to primary
if attempt == max_attempts and can_failover:
if attempt > 1:
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:
uuid.UUID(str(value))
except ValueError:
raise ValidationError("Must be a valid UUID")
cursor.execute(SET_CONFIG_QUERY, [parameter, value])
wrapper_cm = (
conn.execute_wrapper(_query_failover)
if can_failover and alias == replica_alias
else nullcontext()
)
with wrapper_cm:
yielded_cursor = True
yield cursor
_fallback["caller_exited_cleanly"] = True
return
except OperationalError as e:
if yielded_cursor:
if _fallback["succeeded"] and _fallback["caller_exited_cleanly"]:
# Caller's queries succeeded on primary via failover.
# This error is transaction.atomic() cleanup on the
# dead replica connection, suppress it.
return
raise
if not can_failover or attempt == max_attempts:
raise
try:
connections[alias].close()
except Exception:
pass # Best-effort; connection may already be dead
# 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 _fallback["token"] is not None:
reset_read_db_alias(_fallback["token"])
_fallback["token"] = None
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
+38 -2
View File
@@ -1,8 +1,12 @@
"""Tests for rls_transaction retry and fallback logic."""
from unittest.mock import patch
import pytest
from api.db_utils import rls_transaction
from django.db import DEFAULT_DB_ALIAS
from api.db_utils import POSTGRES_TENANT_VAR, rls_transaction
from conftest import TEST_REPLICA_ALIAS
from django.db import DEFAULT_DB_ALIAS, OperationalError, connections
from psycopg2 import OperationalError as Psycopg2OperationalError
from rest_framework_json_api.serializers import ValidationError
@@ -36,3 +40,35 @@ class TestRLSTransaction:
cursor.execute("SELECT current_setting(%s, true)", [custom_param])
result = cursor.fetchone()
assert result == (str(tenant.id),)
@pytest.mark.requires_test_replica_alias
@pytest.mark.django_db(
transaction=True, databases=[DEFAULT_DB_ALIAS, TEST_REPLICA_ALIAS]
)
def test_mid_query_replica_connection_loss_falls_back_to_primary(self, tenant):
"""Real Django connection state: closed replica atomic falls back to primary."""
replica = connections[TEST_REPLICA_ALIAS]
sql = "SELECT current_setting(%s, true), %s"
params = [POSTGRES_TENANT_VAR, 42]
failed_once = {"value": False}
def close_replica_and_raise(execute, sql_arg, params_arg, many, context):
if not failed_once["value"] and sql_arg == sql:
failed_once["value"] = True
replica.close()
try:
raise Psycopg2OperationalError("SSL SYSCALL error: EOF detected")
except Psycopg2OperationalError as psycopg_error:
raise OperationalError(
"SSL SYSCALL error: EOF detected"
) from psycopg_error
return execute(sql_arg, params_arg, many, context)
with patch("api.db_utils.READ_REPLICA_ALIAS", TEST_REPLICA_ALIAS):
with rls_transaction(str(tenant.id), using=TEST_REPLICA_ALIAS) as cursor:
with replica.execute_wrapper(close_replica_and_raise):
cursor.execute(sql, params)
result = cursor.fetchone()
assert failed_once["value"]
assert result == (str(tenant.id), 42)
api/src/backend/api/tests/test_db_utils.py
+573 -19
View File
@@ -1,11 +1,16 @@
from contextlib import contextmanager
from datetime import UTC, datetime
from enum import Enum
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock, call, patch
import pytest
from api.db_utils import (
POSTGRES_TENANT_VAR,
SET_CONFIG_QUERY,
SET_TRANSACTION_READ_ONLY_QUERY,
PostgresEnumMigration,
_is_replica_connection_failure,
_is_safe_primary_replay,
_should_create_index_on_partition,
batch_delete,
create_objects_in_batches,
@@ -392,10 +397,23 @@ class TestRlsTransaction:
with patch("api.db_utils.get_read_db_alias", return_value=None):
with patch("api.db_utils.connections") as mock_connections:
mock_conn = MagicMock()
mock_cursor = MagicMock()
mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
mock_connections.__getitem__.return_value = mock_conn
mock_replica_conn = MagicMock()
mock_replica_cursor = MagicMock()
mock_replica_conn.cursor.return_value.__enter__.return_value = (
mock_replica_cursor
)
mock_primary_conn = MagicMock()
mock_primary_cursor = MagicMock()
mock_primary_conn.cursor.return_value.__enter__.return_value = (
mock_primary_cursor
)
def connections_getitem(alias):
if alias == "replica":
return mock_replica_conn
return mock_primary_conn
mock_connections.__getitem__.side_effect = connections_getitem
mock_connections.__contains__.return_value = True
with patch("api.db_utils.transaction.atomic"):
@@ -525,7 +543,7 @@ class TestRlsTransaction:
def atomic_side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count < 3:
if call_count < 4:
raise OperationalError("Connection error")
return MagicMock(
__enter__=MagicMock(return_value=None),
@@ -544,10 +562,11 @@ class TestRlsTransaction:
with rls_transaction(tenant_id):
pass
assert mock_sleep.call_count == 2
assert mock_sleep.call_count == 3
mock_sleep.assert_any_call(0.5)
mock_sleep.assert_any_call(1.0)
assert mock_logger.info.call_count == 2
mock_sleep.assert_any_call(2.0)
assert mock_logger.info.call_count == 3
def test_rls_transaction_operational_error_inside_context_no_retry(
self, tenants_fixture, enable_read_replica
@@ -578,11 +597,12 @@ class TestRlsTransaction:
raise OperationalError("Conflict with recovery")
mock_sleep.assert_not_called()
mock_conn.close.assert_not_called()
def test_rls_transaction_max_three_attempts_for_replica(
def test_rls_transaction_max_attempts_for_replica(
self, tenants_fixture, enable_read_replica
):
"""Test maximum 3 attempts for replica database."""
"""Test REPLICA_MAX_ATTEMPTS replica tries + 1 primary fallback."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
@@ -606,7 +626,11 @@ class TestRlsTransaction:
with rls_transaction(tenant_id):
pass
assert mock_atomic.call_count == 3
assert mock_atomic.call_args_list[-1] == call(
using=DEFAULT_DB_ALIAS
)
# 3 replica + 1 primary = 4 total
assert mock_atomic.call_count == 4
def test_rls_transaction_replica_no_retry_when_disabled(
self, tenants_fixture, enable_read_replica
@@ -617,10 +641,23 @@ class TestRlsTransaction:
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
mock_conn = MagicMock()
mock_cursor = MagicMock()
mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
mock_connections.__getitem__.return_value = mock_conn
mock_replica_conn = MagicMock()
mock_replica_cursor = MagicMock()
mock_replica_conn.cursor.return_value.__enter__.return_value = (
mock_replica_cursor
)
mock_primary_conn = MagicMock()
mock_primary_cursor = MagicMock()
mock_primary_conn.cursor.return_value.__enter__.return_value = (
mock_primary_cursor
)
def connections_getitem(alias):
if alias == "replica":
return mock_replica_conn
return mock_primary_conn
mock_connections.__getitem__.side_effect = connections_getitem
mock_connections.__contains__.return_value = True
with patch("api.db_utils.transaction.atomic") as mock_atomic:
@@ -682,7 +719,7 @@ class TestRlsTransaction:
def atomic_side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count < 3:
if call_count < 4:
raise OperationalError("Replica error")
return MagicMock(
__enter__=MagicMock(return_value=None),
@@ -691,7 +728,7 @@ class TestRlsTransaction:
with patch(
"api.db_utils.transaction.atomic", side_effect=atomic_side_effect
):
) as mock_atomic:
with patch("api.db_utils.time.sleep"):
with patch(
"api.db_utils.set_read_db_alias", return_value="token"
@@ -701,6 +738,9 @@ class TestRlsTransaction:
with rls_transaction(tenant_id):
pass
assert mock_atomic.call_args_list[-1] == call(
using=DEFAULT_DB_ALIAS
)
mock_logger.warning.assert_called_once()
warning_msg = mock_logger.warning.call_args[0][0]
assert "falling back to primary DB" in warning_msg
@@ -725,7 +765,7 @@ class TestRlsTransaction:
def atomic_side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count < 3:
if call_count < 4:
raise OperationalError("Replica error")
return MagicMock(
__enter__=MagicMock(return_value=None),
@@ -744,7 +784,7 @@ class TestRlsTransaction:
with rls_transaction(tenant_id):
pass
assert mock_logger.info.call_count == 2
assert mock_logger.info.call_count == 3
assert mock_logger.warning.call_count == 1
def test_rls_transaction_operational_error_raised_immediately_on_primary(
@@ -910,6 +950,520 @@ class TestRlsTransaction:
result = cursor.fetchone()
assert result[0] == 1
# --- Mid-query failover tests ---
class _FakeDatabaseError(Exception):
def __init__(self, message, pgcode=None):
super().__init__(message)
self.pgcode = pgcode
def _install_execute_wrapper(self, connection):
connection.execute_wrappers = []
@contextmanager
def _execute_wrapper(fn):
connection.execute_wrappers.append(fn)
try:
yield
finally:
connection.execute_wrappers.remove(fn)
connection.execute_wrapper = _execute_wrapper
def _mock_replica_and_primary_connections(self, mock_connections):
mock_replica_conn = MagicMock()
self._install_execute_wrapper(mock_replica_conn)
mock_replica_cursor = MagicMock()
mock_replica_conn.cursor.return_value.__enter__.return_value = (
mock_replica_cursor
)
mock_primary_conn = MagicMock()
mock_primary_cursor = MagicMock()
mock_primary_raw_cursor = MagicMock()
mock_primary_cursor.cursor = mock_primary_raw_cursor
mock_primary_conn.cursor.return_value = mock_primary_cursor
def connections_getitem(alias):
if alias == "replica":
return mock_replica_conn
return mock_primary_conn
mock_connections.__getitem__.side_effect = connections_getitem
mock_connections.__contains__.return_value = True
return mock_replica_conn, mock_primary_conn, mock_primary_cursor
@pytest.mark.parametrize(
"error",
[
_FakeDatabaseError("connection lost", pgcode="08006"),
_FakeDatabaseError("terminating connection", pgcode="57P01"),
OperationalError("SSL SYSCALL error: EOF detected"),
OperationalError("server closed the connection unexpectedly"),
OperationalError("database system is starting up"),
],
)
def test_replica_connection_failure_detection_allows_failover(self, error):
assert _is_replica_connection_failure(error)
@pytest.mark.parametrize(
"error",
[
_FakeDatabaseError("canceling statement", pgcode="57014"),
_FakeDatabaseError("could not serialize access", pgcode="40001"),
_FakeDatabaseError("deadlock detected", pgcode="40P01"),
OperationalError("deadlock detected"),
],
)
def test_replica_connection_failure_detection_rejects_query_errors(self, error):
assert not _is_replica_connection_failure(error)
@pytest.mark.parametrize(
("sql", "many", "expected"),
[
("SELECT 1", False, True),
(" -- leading comment\nSELECT 1", False, True),
("/* leading comment */ SELECT 1", False, True),
("SELECT 1", True, False),
("SELECTING 1", False, False),
("INSERT INTO fake_table (name) VALUES (%s)", False, False),
("WITH rows AS (SELECT 1) SELECT * FROM rows", False, False),
("SELECT * INTO fake_table_copy FROM fake_table", False, False),
("SELECT * FROM fake_table FOR UPDATE", False, False),
("SELECT * FROM fake_table FOR SHARE", False, False),
],
)
def test_primary_replay_safety_detection(self, sql, many, expected):
assert _is_safe_primary_replay(sql, many) is expected
def test_mid_query_failure_falls_directly_back_to_primary(
self, tenants_fixture, enable_read_replica
):
"""Mid-query replica connection loss is replayed once on primary."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
mock_primary_conn,
mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
outer_atomic = MagicMock()
outer_atomic.__enter__ = MagicMock(return_value=None)
outer_atomic.__exit__ = MagicMock(return_value=False)
fallback_atomic = MagicMock()
fallback_atomic.__enter__ = MagicMock(return_value=None)
fallback_atomic.__exit__ = MagicMock(return_value=False)
with patch(
"api.db_utils.transaction.atomic",
side_effect=[outer_atomic, fallback_atomic],
) as mock_atomic:
with patch("api.db_utils.time.sleep") as mock_sleep:
with patch(
"api.db_utils.set_read_db_alias",
side_effect=["replica-token", "primary-token"],
) as mock_set_alias:
with patch(
"api.db_utils.reset_read_db_alias"
) as mock_reset_alias:
with rls_transaction(tenant_id):
wrapper = mock_replica_conn.execute_wrappers[0]
context_cursor = MagicMock()
mock_execute = MagicMock(
side_effect=OperationalError(
"SSL SYSCALL error: EOF detected"
)
)
wrapper(
mock_execute,
"SELECT %s",
["value"],
False,
{"cursor": context_cursor},
)
mock_sleep.assert_not_called()
(
mock_replica_conn.ensure_connection.assert_not_called()
)
mock_replica_conn.close.assert_called_once()
(
mock_primary_conn.ensure_connection.assert_called_once()
)
mock_primary_conn.cursor.assert_called_once_with()
mock_primary_cursor.execute.assert_has_calls(
[
call(SET_TRANSACTION_READ_ONLY_QUERY),
call(
SET_CONFIG_QUERY,
[POSTGRES_TENANT_VAR, tenant_id],
),
call("SELECT %s", ["value"]),
]
)
assert context_cursor.db == mock_primary_conn
assert (
context_cursor.cursor
== mock_primary_cursor.cursor
)
mock_set_alias.assert_has_calls(
[
call(enable_read_replica),
call(DEFAULT_DB_ALIAS),
]
)
mock_reset_alias.assert_has_calls(
[call("primary-token"), call("replica-token")]
)
assert mock_atomic.call_args_list == [
call(using=enable_read_replica),
call(using=DEFAULT_DB_ALIAS),
]
assert mock_replica_conn.execute_wrappers == []
@pytest.mark.parametrize(
("sql", "params", "many"),
[
("INSERT INTO fake_table (name) VALUES (%s)", [("one",), ("two",)], True),
("INSERT INTO fake_table (name) VALUES (%s)", ["one"], False),
("UPDATE fake_table SET name = %s", ["one"], False),
("DELETE FROM fake_table WHERE id = %s", [1], False),
(
"WITH deleted AS (DELETE FROM fake_table RETURNING *) "
"SELECT * FROM deleted",
None,
False,
),
("SELECT * INTO fake_table_copy FROM fake_table", None, False),
],
)
def test_mid_query_fallback_rejects_unsafe_replay(
self, tenants_fixture, enable_read_replica, sql, params, many
):
"""Only single SELECT statements are replayed on primary."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
mock_primary_conn,
mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
with patch("api.db_utils.transaction.atomic") as mock_atomic:
mock_atomic.return_value.__enter__ = MagicMock(return_value=None)
mock_atomic.return_value.__exit__ = MagicMock(return_value=False)
with patch(
"api.db_utils.set_read_db_alias",
side_effect=["replica-token", "primary-token"],
):
with patch("api.db_utils.reset_read_db_alias"):
with rls_transaction(tenant_id):
wrapper = mock_replica_conn.execute_wrappers[0]
mock_execute = MagicMock(
side_effect=OperationalError(
"server closed the connection"
)
)
with pytest.raises(OperationalError):
wrapper(
mock_execute,
sql,
params,
many,
{"cursor": MagicMock()},
)
mock_primary_conn.ensure_connection.assert_not_called()
mock_primary_conn.cursor.assert_not_called()
mock_primary_cursor.execute.assert_not_called()
mock_primary_cursor.executemany.assert_not_called()
def test_mid_query_non_connection_error_does_not_fall_back(
self, tenants_fixture, enable_read_replica
):
"""Query/concurrency errors are not replayed on primary."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
mock_primary_conn,
_mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
with patch("api.db_utils.transaction.atomic") as mock_atomic:
mock_atomic.return_value.__enter__ = MagicMock(return_value=None)
mock_atomic.return_value.__exit__ = MagicMock(return_value=False)
with patch(
"api.db_utils.set_read_db_alias", return_value="replica-token"
):
with patch("api.db_utils.reset_read_db_alias"):
with rls_transaction(tenant_id):
wrapper = mock_replica_conn.execute_wrappers[0]
mock_execute = MagicMock(
side_effect=OperationalError("deadlock detected")
)
with pytest.raises(OperationalError):
wrapper(
mock_execute,
"SELECT 1",
None,
False,
{"cursor": MagicMock()},
)
mock_replica_conn.close.assert_not_called()
(
mock_primary_conn.ensure_connection.assert_not_called()
)
def test_mid_query_primary_replay_failure_propagates(
self, tenants_fixture, enable_read_replica
):
"""Primary fallback errors propagate as Django OperationalError."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
_mock_primary_conn,
mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
mock_primary_cursor.execute.side_effect = [
None,
None,
OperationalError("primary down"),
]
with patch("api.db_utils.transaction.atomic") as mock_atomic:
mock_atomic.return_value.__enter__ = MagicMock(return_value=None)
mock_atomic.return_value.__exit__ = MagicMock(return_value=False)
with patch(
"api.db_utils.set_read_db_alias",
side_effect=["replica-token", "primary-token"],
):
with patch("api.db_utils.reset_read_db_alias"):
with pytest.raises(OperationalError, match="primary down"):
with rls_transaction(tenant_id):
wrapper = mock_replica_conn.execute_wrappers[0]
mock_execute = MagicMock(
side_effect=OperationalError(
"server closed the connection"
)
)
wrapper(
mock_execute,
"SELECT 1",
None,
False,
{"cursor": MagicMock()},
)
mock_primary_cursor.close.assert_called_once()
def test_mid_query_fallback_suppresses_cleanup_error(
self, tenants_fixture, enable_read_replica
):
"""After successful primary fallback, replica cleanup error is suppressed."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
_mock_primary_conn,
mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
# Replica's atomic.__exit__ raises on dead replica cleanup;
# primary's atomic.__exit__ returns False (healthy commit).
mock_outer_atomic = MagicMock()
mock_outer_atomic.__enter__ = MagicMock(return_value=None)
mock_outer_atomic.__exit__ = MagicMock(
side_effect=OperationalError("cleanup failed on dead replica")
)
mock_fallback_atomic = MagicMock()
mock_fallback_atomic.__enter__ = MagicMock(return_value=None)
mock_fallback_atomic.__exit__ = MagicMock(return_value=False)
atomic_call_count = 0
def atomic_side_effect(*args, **kwargs):
nonlocal atomic_call_count
atomic_call_count += 1
if atomic_call_count == 1:
return mock_outer_atomic
return mock_fallback_atomic
with patch(
"api.db_utils.transaction.atomic",
side_effect=atomic_side_effect,
):
with patch(
"api.db_utils.set_read_db_alias",
side_effect=["replica-token", "primary-token"],
):
with patch("api.db_utils.reset_read_db_alias"):
with rls_transaction(tenant_id):
wrapper = mock_replica_conn.execute_wrappers[0]
mock_execute = MagicMock(
side_effect=OperationalError(
"server closed the connection"
)
)
mock_context = {"cursor": MagicMock()}
wrapper(
mock_execute,
"SELECT 1",
None,
False,
mock_context,
)
mock_primary_cursor.execute.assert_has_calls(
[
call(SET_TRANSACTION_READ_ONLY_QUERY),
call(
SET_CONFIG_QUERY,
[POSTGRES_TENANT_VAR, tenant_id],
),
call("SELECT 1", None),
]
)
def test_wrapper_not_installed_on_primary(self, tenants_fixture):
"""execute_wrapper is not installed when targeting primary DB."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=None):
with patch("api.db_utils.connections") as mock_connections:
mock_conn = MagicMock()
mock_conn.execute_wrappers = []
mock_cursor = MagicMock()
mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
mock_connections.__getitem__.return_value = mock_conn
mock_connections.__contains__.return_value = True
with patch("api.db_utils.transaction.atomic") as mock_atomic:
mock_atomic.return_value.__enter__ = MagicMock(return_value=None)
mock_atomic.return_value.__exit__ = MagicMock(return_value=False)
with rls_transaction(tenant_id):
# No wrapper installed on primary
assert len(mock_conn.execute_wrappers) == 0
def test_stale_connection_closed_on_pre_yield_retry(
self, tenants_fixture, enable_read_replica
):
"""Stale connection is closed before each pre-yield retry."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
mock_conn = MagicMock()
mock_conn.execute_wrappers = []
mock_cursor = MagicMock()
mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
mock_connections.__getitem__.return_value = mock_conn
mock_connections.__contains__.return_value = True
call_count = 0
def atomic_side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count < 3:
raise OperationalError("Connection error")
return MagicMock(
__enter__=MagicMock(return_value=None),
__exit__=MagicMock(return_value=False),
)
with patch(
"api.db_utils.transaction.atomic", side_effect=atomic_side_effect
):
with patch("api.db_utils.time.sleep"):
with patch(
"api.db_utils.set_read_db_alias", return_value="token"
):
with patch("api.db_utils.reset_read_db_alias"):
with rls_transaction(tenant_id):
pass
# close() called for each failed pre-yield attempt
assert mock_conn.close.call_count == 2
def test_caller_error_propagates_after_successful_failover(
self, tenants_fixture, enable_read_replica
):
"""OperationalError raised by caller after failover is NOT suppressed."""
tenant = tenants_fixture[0]
tenant_id = str(tenant.id)
with patch("api.db_utils.get_read_db_alias", return_value=enable_read_replica):
with patch("api.db_utils.connections") as mock_connections:
(
mock_replica_conn,
_mock_primary_conn,
_mock_primary_cursor,
) = self._mock_replica_and_primary_connections(mock_connections)
# Transaction cleanup succeeds so the caller error should surface.
mock_atomic_cm = MagicMock()
mock_atomic_cm.__enter__ = MagicMock(return_value=None)
mock_atomic_cm.__exit__ = MagicMock(return_value=False)
with patch(
"api.db_utils.transaction.atomic", return_value=mock_atomic_cm
):
with patch("api.db_utils.time.sleep"):
with patch(
"api.db_utils.set_read_db_alias", return_value="token"
):
with patch("api.db_utils.reset_read_db_alias"):
with pytest.raises(
OperationalError, match="caller error"
):
with rls_transaction(tenant_id):
# Trigger failover (succeeds on primary)
wrapper = mock_replica_conn.execute_wrappers[0]
mock_execute = MagicMock(
side_effect=OperationalError(
"server closed the connection"
)
)
mock_context = {"cursor": MagicMock()}
wrapper(
mock_execute,
"SELECT 1",
None,
False,
mock_context,
)
# Caller errors after successful failover
# should still propagate.
raise OperationalError("caller error")
class TestPostgresEnumMigration:
"""
api/src/backend/conftest.py
+20
View File
@@ -70,6 +70,7 @@ API_JSON_CONTENT_TYPE = "application/vnd.api+json"
NO_TENANT_HTTP_STATUS = status.HTTP_401_UNAUTHORIZED
TEST_USER = "dev@prowler.com"
TEST_PASSWORD = "testing_psswd"
TEST_REPLICA_ALIAS = "test_replica"
def _install_compliance_catalog_test_cache() -> None:
@@ -2542,8 +2543,27 @@ def pytest_collection_modifyitems(items):
"""Ensure test_rbac.py is executed first."""
items.sort(key=lambda item: 0 if "test_rbac.py" in item.nodeid else 1)
if any(item.get_closest_marker("requires_test_replica_alias") for item in items):
default_database = settings.DATABASES["default"]
if TEST_REPLICA_ALIAS not in settings.DATABASES:
settings.DATABASES[TEST_REPLICA_ALIAS] = {
**default_database,
"TEST": {
**default_database.get("TEST", {}),
"MIRROR": "default",
},
}
django_connections.databases[TEST_REPLICA_ALIAS] = settings.DATABASES[
TEST_REPLICA_ALIAS
]
def pytest_configure(config):
config.addinivalue_line(
"markers",
"requires_test_replica_alias: creates a test-only replica alias mirrored "
"to default",
)
# Apply the mock before the test session starts. This is necessary to avoid admin error when running the
# 0004_rbac_missing_admin_roles migration
patch("api.db_router.MainRouter.admin_db", new="default").start()