Allow for explicit use of database replicas

Author: emyller

settings/dev.py

@@ -33,7 +33,12 @@
     "task_processor": dj_database_url.parse(
         "postgresql://postgres@localhost:6544/postgres",
     ),
+    # Dummy replicas
+    "replica_1": default_database_url,
+    "replica_2": default_database_url,
+    "replica_3": default_database_url,
 }
+REPLICA_READ_STRATEGY = "distributed"
 TASK_PROCESSOR_DATABASES = ["default"]
 INSTALLED_APPS = [
     "django.contrib.auth",

src/common/core/utils.py

@@ -1,16 +1,27 @@
+import enum
 import json
 import pathlib
+import random
 from functools import lru_cache
-from typing import NotRequired, TypedDict
+from itertools import cycle
+from typing import NotRequired, TypedDict, TypeVar
 
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import AbstractBaseUser
-from django.db.models import Manager
+from django.db import connections
+from django.db.models import Manager, Model
 
 UNKNOWN = "unknown"
 VERSIONS_INFO_FILE_LOCATION = ".versions.json"
 
+ModelType = TypeVar("ModelType", bound=Model)
+
+
+class ReplicaReadStrategy(enum.StrEnum):
+    DISTRIBUTED = enum.auto()
+    SEQUENTIAL = enum.auto()
+
 
 class SelfHostedData(TypedDict):
     has_users: bool
@@ -114,3 +125,24 @@ def get_file_contents(file_path: str) -> str | None:
             return f.read().replace("\n", "")
     except FileNotFoundError:
         return None
+
+
+def using_database_replica(manager: Manager[ModelType]) -> Manager[ModelType]:
+    local_replicas = [name for name in connections if name.startswith("replica_")]
+
+    if not local_replicas:
+        return manager
+
+    if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.SEQUENTIAL:
+        global _sequential_replica_manager
+        try:
+            isinstance(_sequential_replica_manager, cycle)  # type: ignore[name-defined]
+        except NameError:
+            _sequential_replica_manager = cycle(local_replicas)  # type: ignore[name-defined]
+        finally:
+            return manager.db_manager(next(_sequential_replica_manager))  # type: ignore[name-defined]
+
+    if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.DISTRIBUTED:
+        return manager.db_manager(random.choice(local_replicas))
+
+    raise NotImplementedError(f"Unsupported strategy: {settings.REPLICA_READ_STRATEGY}")

tests/unit/common/core/test_utils.py

@@ -1,8 +1,10 @@
 import json
 
 import pytest
+from django.contrib.auth import get_user_model
 from pyfakefs.fake_filesystem import FakeFilesystem
-from pytest_django.fixtures import SettingsWrapper
+from pytest_django.fixtures import DjangoAssertNumQueries, SettingsWrapper
+from pytest_mock import MockerFixture
 
 from common.core.utils import (
     get_file_contents,
@@ -13,6 +15,7 @@
     is_enterprise,
     is_oss,
     is_saas,
+    using_database_replica,
 )
 
 pytestmark = pytest.mark.django_db
@@ -181,3 +184,85 @@ def test_get_version__invalid_file_contents__returns_unknown(
 
     # Then
     assert result == "unknown"
+
+
+@pytest.mark.django_db(databases="__all__")
+def test_maybe_use_replicas__no_replicas__points_to_default(
+    django_assert_num_queries: DjangoAssertNumQueries,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch("common.core.utils.connections", {"default": {}})
+    manager = get_user_model().objects
+
+    # When / Then
+    with django_assert_num_queries(1, using="default"):
+        using_database_replica(manager).first()
+
+
+@pytest.mark.django_db(databases="__all__")
+def test_maybe_use_replicas__with_replicas__sequential_strategy__picks_databases_sequentially(
+    django_assert_num_queries: DjangoAssertNumQueries,
+    settings: SettingsWrapper,
+) -> None:
+    # Given
+    settings.REPLICA_READ_STRATEGY = "sequential"
+    manager = get_user_model().objects
+
+    # When / Then
+    with django_assert_num_queries(1, using="replica_1"):
+        using_database_replica(manager).first()
+    with django_assert_num_queries(1, using="replica_2"):
+        using_database_replica(manager).first()
+    with django_assert_num_queries(1, using="replica_3"):
+        using_database_replica(manager).first()
+    with django_assert_num_queries(1, using="replica_1"):
+        using_database_replica(manager).first()
+
+
+@pytest.mark.django_db(databases="__all__")
+def test_maybe_use_replicas__with_replicas__distributed_strategy__picks_databases_randomly(
+    django_assert_max_num_queries: DjangoAssertNumQueries,
+    settings: SettingsWrapper,
+) -> None:
+    # Given
+    settings.REPLICA_READ_STRATEGY = "distributed"
+    manager = get_user_model().objects
+
+    # When / Then
+    with django_assert_max_num_queries(10, using="replica_1") as captured:
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+    assert (captured.final_queries or 0) >= 1
+    with django_assert_max_num_queries(10, using="replica_2") as captured:
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+    assert (captured.final_queries or 0) >= 1
+    with django_assert_max_num_queries(10, using="replica_3") as captured:
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+        using_database_replica(manager).first()
+    assert (captured.final_queries or 0) >= 1