Skip to content

Commit 658d14d

Browse files
michael.yakmichaelyaakoby
authored andcommitted
Suport multiple DB connections by naming respective health-provider and composing them
Issue #96 asked to support for showing health of multiple DB engines/endpoints. To address this, Pyctuator now supports `ComposeiteHealthProvider` which wraps around a list of health-providers. Also, the built-in health-providers now support overrding the default name. Setting health check for multiple DB engines can be done as follows: ```python pyctuator.register_health_provider( CompositeHealthProvider( "db", DbHealthProvider(db_engine, "db1"), DbHealthProvider(db_engine, "db2"), ) ) ```
1 parent 995413d commit 658d14d

File tree

6 files changed

+136
-6
lines changed

6 files changed

+136
-6
lines changed
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
from dataclasses import dataclass
2+
from typing import Mapping
3+
4+
from pyctuator.health.health_provider import HealthProvider, HealthStatus, Status
5+
6+
7+
@dataclass
8+
class CompositeHealthStatus(HealthStatus):
9+
status: Status
10+
details: Mapping[str, HealthStatus] # type: ignore[assignment]
11+
12+
13+
class CompositeHealthProvider(HealthProvider):
14+
15+
def __init__(self, name: str, *health_providers: HealthProvider) -> None:
16+
super().__init__()
17+
self.name = name
18+
self.health_providers = health_providers
19+
20+
def is_supported(self) -> bool:
21+
return True
22+
23+
def get_name(self) -> str:
24+
return self.name
25+
26+
def get_health(self) -> CompositeHealthStatus:
27+
health_statuses: Mapping[str, HealthStatus] = {
28+
provider.get_name(): provider.get_health()
29+
for provider in self.health_providers
30+
if provider.is_supported()
31+
}
32+
33+
# Health is UP if no provider is registered
34+
if not health_statuses:
35+
return CompositeHealthStatus(Status.UP, health_statuses)
36+
37+
# If there's at least one provider and any of the providers is DOWN, the service is DOWN
38+
service_is_down = any(health_status.status == Status.DOWN for health_status in health_statuses.values())
39+
if service_is_down:
40+
return CompositeHealthStatus(Status.DOWN, health_statuses)
41+
42+
# If there's at least one provider and none of the providers is DOWN and at least one is UP, the service is UP
43+
service_is_up = any(health_status.status == Status.UP for health_status in health_statuses.values())
44+
if service_is_up:
45+
return CompositeHealthStatus(Status.UP, health_statuses)
46+
47+
# else, all providers are unknown so the service is UNKNOWN
48+
return CompositeHealthStatus(Status.UNKNOWN, health_statuses)

pyctuator/health/db_health_provider.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,25 @@ class DbHealthStatus(HealthStatus):
2121

2222
class DbHealthProvider(HealthProvider):
2323

24-
def __init__(self, engine: Engine) -> None:
24+
def __init__(self, engine: Engine, name: str = "db") -> None:
2525
super().__init__()
2626
self.engine = engine
27+
self.name = name
2728

2829
def is_supported(self) -> bool:
2930
return importlib.util.find_spec("sqlalchemy") is not None
3031

3132
def get_name(self) -> str:
32-
return "db"
33+
return self.name
3334

3435
def get_health(self) -> DbHealthStatus:
3536
try:
3637
with self.engine.connect() as conn:
3738
if self.engine.dialect.do_ping(conn.connection): # type: ignore[arg-type]
38-
return DbHealthStatus(status=Status.UP, details=DbHealthDetails(self.engine.name))
39+
return DbHealthStatus(
40+
status=Status.UP,
41+
details=DbHealthDetails(self.engine.name)
42+
)
3943

4044
return DbHealthStatus(
4145
status=Status.UNKNOWN,

pyctuator/health/redis_health_provider.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,16 @@ class RedisHealthStatus(HealthStatus):
2222

2323
class RedisHealthProvider(HealthProvider):
2424

25-
def __init__(self, redis: Redis) -> None:
25+
def __init__(self, redis: Redis, name: str = "redis") -> None:
2626
super().__init__()
2727
self.redis = redis
28+
self.name = name
2829

2930
def is_supported(self) -> bool:
3031
return importlib.util.find_spec("redis") is not None
3132

3233
def get_name(self) -> str:
33-
return "redis"
34+
return self.name
3435

3536
def get_health(self) -> RedisHealthStatus:
3637
try:

pyctuator/impl/pyctuator_impl.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ def get_health(self) -> HealthSummary:
119119
if service_is_down:
120120
return HealthSummary(Status.DOWN, health_statuses)
121121

122-
# IF there's at least one provider and none of the providers is DOWN and at least one is UP, the service is UP
122+
# If there's at least one provider and none of the providers is DOWN and at least one is UP, the service is UP
123123
service_is_up = any(health_status.status == Status.UP for health_status in health_statuses.values())
124124
if service_is_up:
125125
return HealthSummary(Status.UP, health_statuses)
Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
from dataclasses import dataclass
2+
3+
from pyctuator.health.composite_health_provider import CompositeHealthProvider, CompositeHealthStatus
4+
from pyctuator.health.health_provider import HealthProvider, HealthStatus, Status, HealthDetails
5+
6+
7+
@dataclass
8+
class CustomHealthDetails(HealthDetails):
9+
details: str
10+
11+
12+
class CustomHealthProvider(HealthProvider):
13+
14+
def __init__(self, name: str, status: HealthStatus) -> None:
15+
super().__init__()
16+
self.name = name
17+
self.status = status
18+
19+
def is_supported(self) -> bool:
20+
return True
21+
22+
def get_name(self) -> str:
23+
return self.name
24+
25+
def get_health(self) -> HealthStatus:
26+
return self.status
27+
28+
29+
def test_composite_health_provider_no_providers() -> None:
30+
health_provider = CompositeHealthProvider(
31+
"comp1",
32+
)
33+
34+
assert health_provider.get_name() == "comp1"
35+
36+
assert health_provider.get_health() == CompositeHealthStatus(
37+
status=Status.UP,
38+
details={}
39+
)
40+
41+
42+
def test_composite_health_provider_all_up() -> None:
43+
health_provider = CompositeHealthProvider(
44+
"comp2",
45+
CustomHealthProvider("hp1", HealthStatus(Status.UP, CustomHealthDetails("d1"))),
46+
CustomHealthProvider("hp2", HealthStatus(Status.UP, CustomHealthDetails("d2"))),
47+
)
48+
49+
assert health_provider.get_name() == "comp2"
50+
51+
assert health_provider.get_health() == CompositeHealthStatus(
52+
status=Status.UP,
53+
details={
54+
"hp1": HealthStatus(Status.UP, CustomHealthDetails("d1")),
55+
"hp2": HealthStatus(Status.UP, CustomHealthDetails("d2")),
56+
}
57+
)
58+
59+
60+
def test_composite_health_provider_one_down() -> None:
61+
health_provider = CompositeHealthProvider(
62+
"comp3",
63+
CustomHealthProvider("hp1", HealthStatus(Status.UP, CustomHealthDetails("d1"))),
64+
CustomHealthProvider("hp2", HealthStatus(Status.DOWN, CustomHealthDetails("d2"))),
65+
)
66+
67+
assert health_provider.get_name() == "comp3"
68+
69+
assert health_provider.get_health() == CompositeHealthStatus(
70+
status=Status.DOWN,
71+
details={
72+
"hp1": HealthStatus(Status.UP, CustomHealthDetails("d1")),
73+
"hp2": HealthStatus(Status.DOWN, CustomHealthDetails("d2")),
74+
}
75+
)

tests/health/test_db_health_provider.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ def test_sqlite_health() -> None:
3333
engine = create_engine("sqlite:///:memory:", echo=True)
3434
health_provider = DbHealthProvider(engine)
3535
assert health_provider.get_health() == DbHealthStatus(status=Status.UP, details=DbHealthDetails("sqlite"))
36+
assert health_provider.get_name() == "db"
37+
assert DbHealthProvider(engine, "kuki").get_name() == "kuki"
3638

3739

3840
@pytest.mark.usefixtures("require_sql_alchemy", "require_pymysql", "require_mysql_server")

0 commit comments

Comments
 (0)