From b44f8ebf5ffb308fe8519cfeafdbbf6f79874b21 Mon Sep 17 00:00:00 2001 From: Josema Camacho Date: Thu, 9 Jul 2026 17:47:03 +0200 Subject: [PATCH] fix(api): add query-level retry with primary fallback to `rls_transaction` via `execute_wrapper` (#10379) Co-authored-by: Pepe Fagoaga --- .../rls-transaction-primary-fallback.fixed.md | 1 + api/src/backend/api/db_utils.py | 323 ++++++++-- .../tests/integration/test_rls_transaction.py | 40 +- api/src/backend/api/tests/test_db_utils.py | 592 +++++++++++++++++- api/src/backend/conftest.py | 20 + 5 files changed, 905 insertions(+), 71 deletions(-) create mode 100644 api/changelog.d/rls-transaction-primary-fallback.fixed.md diff --git a/api/changelog.d/rls-transaction-primary-fallback.fixed.md b/api/changelog.d/rls-transaction-primary-fallback.fixed.md new file mode 100644 index 0000000000..4f1035aa9f --- /dev/null +++ b/api/changelog.d/rls-transaction-primary-fallback.fixed.md @@ -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 diff --git a/api/src/backend/api/db_utils.py b/api/src/backend/api/db_utils.py index 2c378f2ea8..b6d3fdada1 100644 --- a/api/src/backend/api/db_utils.py +++ b/api/src/backend/api/db_utils.py @@ -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): diff --git a/api/src/backend/api/tests/integration/test_rls_transaction.py b/api/src/backend/api/tests/integration/test_rls_transaction.py index bd46871586..ce026d4f58 100644 --- a/api/src/backend/api/tests/integration/test_rls_transaction.py +++ b/api/src/backend/api/tests/integration/test_rls_transaction.py @@ -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) diff --git a/api/src/backend/api/tests/test_db_utils.py b/api/src/backend/api/tests/test_db_utils.py index 06b528b44a..347be2b64d 100644 --- a/api/src/backend/api/tests/test_db_utils.py +++ b/api/src/backend/api/tests/test_db_utils.py @@ -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: """ diff --git a/api/src/backend/conftest.py b/api/src/backend/conftest.py index f4e398d76a..daa0616971 100644 --- a/api/src/backend/conftest.py +++ b/api/src/backend/conftest.py @@ -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()