feat(utils): Allow for explicit use of database replicas (#95)

* Improve database settings

- Prevents conflicts with other systems
- Less unused code, yay

* Allow for explicit use of database replicas

* Fall back to cross-region replicas

* Add logs

* Fix coverage

* Fall back to default database if all replicas fail

* Improve name

* Improve logging

* Improve typing

* Fix member placing

* Fix manager typing

Author: emyller

.github/workflows/python-test.yml

@@ -24,14 +24,14 @@ jobs:
         image: postgres:15.5-alpine
         env:
           POSTGRES_HOST_AUTH_METHOD: trust
-        ports: ['5432:5432']
+        ports: ["6543:5432"]
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
 
       task-processor-database:
         image: postgres:15.5-alpine
         env:
           POSTGRES_HOST_AUTH_METHOD: trust
-        ports: ['5433:5432']
+        ports: ["6544:5432"]
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
 
     steps:

docker/docker-compose.local.yml

@@ -16,7 +16,7 @@ services:
     volumes:
       - default-database:/var/lib/postgresql/data
     ports:
-      - 5432:5432
+      - 6543:5432
     environment:
       POSTGRES_HOST_AUTH_METHOD: trust
       TZ: UTC
@@ -31,7 +31,7 @@ services:
     volumes:
       - task-processor-database:/var/lib/postgresql/data
     ports:
-      - 5433:5432
+      - 6544:5432
     environment:
       POSTGRES_HOST_AUTH_METHOD: trust
       TZ: UTC

settings/dev.py

@@ -25,19 +25,22 @@
 # Settings required for tests
 SECRET_KEY = "test"
 DATABASES = {
-    "default": dj_database_url.parse(
-        env(
-            "DATABASE_URL",
-            default="postgresql://postgres@localhost:5432/postgres",
-        ),
+    "default": (
+        default_database_url := dj_database_url.parse(
+            "postgresql://postgres@localhost:6543/postgres",
+        )
     ),
     "task_processor": dj_database_url.parse(
-        env(
-            "TASK_PROCESSOR_DATABASE_URL",
-            default="postgresql://postgres@localhost:5433/postgres",
-        ),
+        "postgresql://postgres@localhost:6544/postgres",
     ),
+    # Dummy replicas
+    "replica_1": default_database_url,
+    "replica_2": default_database_url,
+    "replica_3": default_database_url,
+    "cross_region_replica_1": default_database_url,
+    "cross_region_replica_2": default_database_url,
 }
+REPLICA_READ_STRATEGY = "distributed"
 TASK_PROCESSOR_DATABASES = ["default"]
 INSTALLED_APPS = [
     "django.contrib.auth",

src/common/core/__init__.py

@@ -0,0 +1,6 @@
+import enum
+
+
+class ReplicaReadStrategy(enum.StrEnum):
+    DISTRIBUTED = enum.auto()
+    SEQUENTIAL = enum.auto()

src/common/core/utils.py

@@ -1,16 +1,30 @@
 import json
+import logging
 import pathlib
+import random
 from functools import lru_cache
-from typing import NotRequired, TypedDict
+from itertools import cycle
+from typing import Iterator, Literal, 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
+from django.db.utils import OperationalError
+
+from common.core import ReplicaReadStrategy
+
+logger = logging.getLogger(__name__)
 
 UNKNOWN = "unknown"
 VERSIONS_INFO_FILE_LOCATION = ".versions.json"
 
+ManagerType = TypeVar("ManagerType", bound=Manager[Model])
+
+ReplicaNamePrefix = Literal["replica_", "cross_region_replica_"]
+_replica_sequential_names_by_prefix: dict[ReplicaNamePrefix, Iterator[str]] = {}
+
 
 class SelfHostedData(TypedDict):
     has_users: bool
@@ -114,3 +128,53 @@ def get_file_contents(file_path: str) -> str | None:
             return f.read().replace("\n", "")
     except FileNotFoundError:
         return None
+
+
+def using_database_replica(
+    manager: ManagerType,
+    replica_prefix: ReplicaNamePrefix = "replica_",
+) -> ManagerType:
+    """Attempts to bind a manager to a healthy database replica"""
+    local_replicas = [name for name in connections if name.startswith(replica_prefix)]
+
+    if not local_replicas:
+        logger.info("No replicas set up.")
+        return manager
+
+    chosen_replica = None
+
+    if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.SEQUENTIAL:
+        sequence = _replica_sequential_names_by_prefix.setdefault(
+            replica_prefix, cycle(local_replicas)
+        )
+        for _ in range(len(local_replicas)):
+            attempted_replica = next(sequence)
+            try:
+                connections[attempted_replica].ensure_connection()
+                chosen_replica = attempted_replica
+                break
+            except OperationalError:
+                logger.exception(f"Replica '{attempted_replica}' is not available.")
+                continue
+
+    if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.DISTRIBUTED:
+        for _ in range(len(local_replicas)):
+            attempted_replica = random.choice(local_replicas)
+            try:
+                connections[attempted_replica].ensure_connection()
+                chosen_replica = attempted_replica
+                break
+            except OperationalError:
+                logger.exception(f"Replica '{attempted_replica}' is not available.")
+                local_replicas.remove(attempted_replica)
+                continue
+
+    if not chosen_replica:
+        if replica_prefix == "replica_":
+            logger.warning("Falling back to cross-region replicas, if any.")
+            return using_database_replica(manager, "cross_region_replica_")
+
+        logger.warning("No replicas available.")
+        return manager
+
+    return manager.db_manager(chosen_replica)

tests/__init__.py

@@ -0,0 +1,3 @@
+from typing import Callable
+
+GetLogsFixture = Callable[[str], list[tuple[str, str]]]

tests/conftest.py

@@ -5,9 +5,21 @@
 import prometheus_client
 import pytest
 
+from tests import GetLogsFixture
+
 pytest_plugins = "flagsmith-test-tools"
 
 
+@pytest.fixture()
+def get_logs(caplog: pytest.LogCaptureFixture) -> GetLogsFixture:
+    caplog.set_level("DEBUG")
+
+    def _logs(module: str) -> list[tuple[str, str]]:
+        return [(r.levelname, r.message) for r in caplog.records if r.name == module]
+
+    return _logs
+
+
 @pytest.fixture()
 def prometheus_multiproc_dir(
     tmp_path_factory: pytest.TempPathFactory,