Skip to content

Commit 1c6c522

Browse files
authored
feat(utils): New is_database_replica_setup utility (#100)
* Tell whether any database replica is setup * Cache
1 parent a21d125 commit 1c6c522

File tree

2 files changed

+41
-1
lines changed

2 files changed

+41
-1
lines changed

src/common/core/utils.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import random
55
from functools import lru_cache
66
from itertools import cycle
7-
from typing import Iterator, Literal, NotRequired, TypedDict, TypeVar
7+
from typing import Iterator, Literal, NotRequired, TypedDict, TypeVar, get_args
88

99
from django.conf import settings
1010
from django.contrib.auth import get_user_model
@@ -130,6 +130,14 @@ def get_file_contents(file_path: str) -> str | None:
130130
return None
131131

132132

133+
@lru_cache()
134+
def is_database_replica_setup() -> bool:
135+
"""Checks if any database replica is set up"""
136+
return any(
137+
name for name in connections if name.startswith(get_args(ReplicaNamePrefix))
138+
)
139+
140+
133141
def using_database_replica(
134142
manager: ManagerType,
135143
replica_prefix: ReplicaNamePrefix = "replica_",

tests/unit/common/core/test_utils.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
get_version_info,
1616
get_versions_from_manifest,
1717
has_email_provider,
18+
is_database_replica_setup,
1819
is_enterprise,
1920
is_oss,
2021
is_saas,
@@ -204,6 +205,37 @@ def test_get_version__invalid_file_contents__returns_unknown(
204205
assert result == "unknown"
205206

206207

208+
@pytest.mark.parametrize(
209+
["database_names", "expected"],
210+
[
211+
({"default"}, False),
212+
({"default", "another_database_with_'replica'_in_its_name"}, False),
213+
({"default", "task_processor"}, False),
214+
({"default", "replica_1"}, True),
215+
({"default", "replica_1", "replica_2"}, True),
216+
({"default", "cross_region_replica_1"}, True),
217+
({"default", "replica_1", "cross_region_replica_1"}, True),
218+
],
219+
)
220+
def test_is_database_replica_setup__tells_whether_any_replica_is_present(
221+
database_names: list[str],
222+
expected: bool,
223+
mocker: MockerFixture,
224+
) -> None:
225+
# Given
226+
is_database_replica_setup.cache_clear()
227+
mocker.patch(
228+
"common.core.utils.connections",
229+
{name: connections["default"] for name in database_names},
230+
)
231+
232+
# When
233+
result = is_database_replica_setup()
234+
235+
# Then
236+
assert result is expected
237+
238+
207239
@pytest.mark.django_db(databases="__all__")
208240
def test_using_database_replica__no_replicas__points_to_default(
209241
django_assert_num_queries: DjangoAssertNumQueries,

0 commit comments

Comments
 (0)