fix(db): Fix read replica strategy (#3426)

Flagsmith · Feb 15, 2024 · d63a289 · d63a289
1 parent 16a2468
commit d63a289
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/api/app/routers.py b/api/app/routers.py
@@ -100,9 +100,9 @@ def allow_migrate(self, db, app_label, model_name=None, **hints):
 
     def _get_replica(self, replicas: list[str]) -> None | str:
         while replicas:
-            if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.DISTRIBUTED.value:
+            if settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.DISTRIBUTED:
                 database = random.choice(replicas)
-            elif settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.SEQUENTIAL.value:
+            elif settings.REPLICA_READ_STRATEGY == ReplicaReadStrategy.SEQUENTIAL:
                 database = replicas[0]
             else:
                 raise ImproperlyConfiguredError(

diff --git a/api/tests/unit/app/test_unit_app_routers.py b/api/tests/unit/app/test_unit_app_routers.py
@@ -28,7 +28,7 @@ def test_replica_router_db_for_read_with_one_offline_replica(
 
     # Set unused cross regional db for testing non-inclusion.
     settings.NUM_CROSS_REGION_DB_REPLICAS = 2
-    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED.value
+    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED
 
     conn_patch = mocker.MagicMock()
     conn_patch.is_usable.side_effect = (False, True)
@@ -63,7 +63,7 @@ def test_replica_router_db_for_read_with_local_offline_replicas(
 
     # Use cross regional db for fallback after replicas.
     settings.NUM_CROSS_REGION_DB_REPLICAS = 2
-    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED.value
+    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED
 
     conn_patch = mocker.MagicMock()
 
@@ -107,7 +107,7 @@ def test_replica_router_db_for_read_with_all_offline_replicas(
     # Given
     settings.NUM_DB_REPLICAS = 4
     settings.NUM_CROSS_REGION_DB_REPLICAS = 2
-    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED.value
+    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.DISTRIBUTED
 
     conn_patch = mocker.MagicMock()
 
@@ -141,7 +141,7 @@ def test_replica_router_db_with_sequential_read(
     # Given
     settings.NUM_DB_REPLICAS = 100
     settings.NUM_CROSS_REGION_DB_REPLICAS = 2
-    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.SEQUENTIAL.value
+    settings.REPLICA_READ_STRATEGY = ReplicaReadStrategy.SEQUENTIAL
 
     conn_patch = mocker.MagicMock()