flagsmith-common 2.0.0__tar.gz → 2.2.0__tar.gz

{flagsmith_common-2.0.0 → flagsmith_common-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flagsmith-common
-Version: 2.0.0
+Version: 2.2.0
 Summary: Flagsmith's common library
 License: BSD-3-Clause
 Author: Matthew Elwell

{flagsmith_common-2.0.0 → flagsmith_common-2.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "flagsmith-common"
-version = "2.0.0"
+version = "2.2.0"
 description = "Flagsmith's common library"
 requires-python = ">=3.11,<4.0"
 dependencies = [

flagsmith_common-2.2.0/src/common/core/__init__.py

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

{flagsmith_common-2.0.0 → flagsmith_common-2.2.0}/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, get_args
 
 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,61 @@ def get_file_contents(file_path: str) -> str | None:
             return f.read().replace("\n", "")
     except FileNotFoundError:
         return None
+
+
+@lru_cache()
+def is_database_replica_setup() -> bool:
+    """Checks if any database replica is set up"""
+    return any(
+        name for name in connections if name.startswith(get_args(ReplicaNamePrefix))
+    )
+
+
+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)