Make read/write DB routing consistent for http requests and background tasks (baserow#4077)

Make  read/write DB routing consistent for http requests and background tasks

Author: bram2w

backend/pytest.ini

@@ -50,4 +50,5 @@ markers =
     websockets: All tests related to handeling web socket connections
     import_export_workspace: All tests related to importing and exporting workspaces
     data_sync: All tests related to data sync functionality
+    replica: All tests related to db replicas
     workspace_search: All tests related to workspace search functionality

backend/src/baserow/config/db_routers.py

@@ -1,28 +1,68 @@
 import random
 
 from django.conf import settings
+from django.db import transaction
 
 from asgiref.local import Local
 
-DATABASE_READ_REPLICAS = settings.DATABASE_READ_REPLICAS
 DEFAULT_DB_ALIAS = "default"
 
 _db_state = Local()
 
 
-def set_write_mode():
-    _db_state.pinned = True
+def set_db_alias(alias: str) -> str:
+    """
+    Pin the current db connection alias to use for the current request or celery task.
+
+    :param alias: The database alias to pin to.
+    :return: The database alias that was set.
+    """
+
+    _db_state.alias = alias
+    return alias
+
+
+def get_db_alias() -> str | None:
+    """
+    Get the pinned db connection alias for the current request or celery task.
+    """
+
+    return getattr(_db_state, "alias", None)
+
+
+def set_db_alias_for_read():
+    """
+    Choose a read replica for read queries, unless we are in an atomic block,
+    in which case we should use the primary database to avoid replication lag issues
+    or trying to lock data in a read replica.
+    Once a read replica is chosen, it is pinned for the duration of the request or
+    celery task.
+
+    :return: The database alias to use for read queries.
+    """
+
+    # Make sure LOCK always happen on the DEFAULT_DB_ALIAS
+    if transaction.get_connection().in_atomic_block:
+        return set_db_alias(DEFAULT_DB_ALIAS)
+
+    # If we already have an alias set, return it
+    if (alias := get_db_alias()) is not None:
+        return alias
 
+    # Choose a random read replica if available, otherwise use the default
+    if settings.DATABASE_READ_REPLICAS:
+        alias = random.choice(settings.DATABASE_READ_REPLICAS)  # nosec
+    else:
+        alias = DEFAULT_DB_ALIAS
 
-def is_write_mode():
-    return getattr(_db_state, "pinned", False)
+    return set_db_alias(alias)
 
 
 def clear_db_state():
     """Should be called when a request or celery finishes."""
 
-    if hasattr(_db_state, "pinned"):
-        del _db_state.pinned
+    if hasattr(_db_state, "alias"):
+        del _db_state.alias
 
 
 class ReadReplicaRouter:
@@ -35,20 +75,14 @@ class ReadReplicaRouter:
     """
 
     def db_for_read(self, model, **hints):
-        if is_write_mode():
-            return DEFAULT_DB_ALIAS
-        if DATABASE_READ_REPLICAS:
-            read = random.choice(DATABASE_READ_REPLICAS)  # nosec
-            return read
-        return DEFAULT_DB_ALIAS
+        return set_db_alias_for_read()
 
     def db_for_write(self, model, **hints):
-        set_write_mode()
-        return DEFAULT_DB_ALIAS
+        return set_db_alias(DEFAULT_DB_ALIAS)
 
     def allow_relation(self, obj1, obj2, **hints):
         db_set = {DEFAULT_DB_ALIAS}
-        db_set.update(DATABASE_READ_REPLICAS)
+        db_set.update(settings.DATABASE_READ_REPLICAS)
         if obj1._state.db in db_set and obj2._state.db in db_set:
             return True
         return None

backend/tests/baserow/config/test_read_replica_router.py

@@ -0,0 +1,95 @@
+from unittest.mock import MagicMock, patch
+
+from django.test.utils import override_settings
+
+import pytest
+
+from baserow.config.db_routers import (
+    DEFAULT_DB_ALIAS,
+    ReadReplicaRouter,
+    clear_db_state,
+    get_db_alias,
+)
+
+
+@pytest.mark.replica
+@pytest.mark.django_db
+@override_settings(DATABASE_READ_REPLICAS=["replica1", "replica2"])
+def test_router_uses_replica_outside_transaction_and_sticks():
+    router = ReadReplicaRouter()
+    clear_db_state()
+
+    mock_conn = MagicMock()
+    mock_conn.get_autocommit.return_value = True
+    mock_conn.in_atomic_block = False
+
+    with patch("django.db.transaction.get_connection", return_value=mock_conn):
+        # Outside transaction should use replica and stick to it
+        alias1 = router.db_for_read(model=None)
+        assert alias1 in ["replica1", "replica2"]
+        assert get_db_alias() == alias1
+
+        alias2 = router.db_for_read(model=None)
+        assert alias2 == alias1
+
+
+@pytest.mark.replica
+@pytest.mark.django_db
+@override_settings(DATABASE_READ_REPLICAS=["replica1", "replica2"])
+def test_router_switches_to_default_inside_transaction():
+    router = ReadReplicaRouter()
+    clear_db_state()
+
+    # Mock connection for outside transaction
+    mock_conn_outside = MagicMock()
+    mock_conn_outside.get_autocommit.return_value = True
+    mock_conn_outside.in_atomic_block = False
+
+    with patch("django.db.transaction.get_connection", return_value=mock_conn_outside):
+        # Start outside transaction - should use replica
+        alias_outside = router.db_for_read(model=None)
+        assert alias_outside in ["replica1", "replica2"]
+
+    # Mock connection for inside transaction
+    mock_conn_inside = MagicMock()
+    mock_conn_inside.get_autocommit.return_value = False
+    mock_conn_inside.in_atomic_block = True
+
+    with patch("django.db.transaction.get_connection", return_value=mock_conn_inside):
+        # Enter transaction - should switch to default and stick
+        alias1 = router.db_for_read(model=None)
+        assert alias1 == DEFAULT_DB_ALIAS
+
+        alias2 = router.db_for_read(model=None)
+        assert alias2 == DEFAULT_DB_ALIAS
+
+    # After transaction, should still be default (sticky)
+    with patch("django.db.transaction.get_connection", return_value=mock_conn_outside):
+        alias3 = router.db_for_read(model=None)
+        assert alias3 == DEFAULT_DB_ALIAS
+
+
+@pytest.mark.replica
+@pytest.mark.django_db
+@override_settings(DATABASE_READ_REPLICAS=["replica1", "replica2"])
+def test_write_switches_to_default_and_sticks():
+    router = ReadReplicaRouter()
+    clear_db_state()
+
+    # Mock connection for outside transaction
+    mock_conn = MagicMock()
+    mock_conn.get_autocommit.return_value = True
+    mock_conn.in_atomic_block = False
+
+    with patch("django.db.transaction.get_connection", return_value=mock_conn):
+        # Start outside transaction - should use replica
+        alias_before_write = router.db_for_read(model=None)
+        assert alias_before_write in ["replica1", "replica2"]
+
+        # Write should switch to default
+        write_alias = router.db_for_write(model=None)
+        assert write_alias == DEFAULT_DB_ALIAS
+
+        # Read after write should still be default
+        read_after_write = router.db_for_read(model=None)
+        assert read_after_write == DEFAULT_DB_ALIAS

changelog/entries/unreleased/refactor/3848_make_readwrite_db_routing_consistent_for_http_requests_and_b.json

@@ -0,0 +1,8 @@
+{
+    "type": "refactor",
+    "message": "Make read/write DB routing consistent for http requests and background tasks",
+    "domain": "database",
+    "issue_number": 3848,
+    "bullet_points": [],
+    "created_at": "2025-10-06"
+}