Skip to content

Commit

Permalink
fix health check and skip not needed tasks (reflex-dev#4563)
Browse files Browse the repository at this point in the history
  • Loading branch information
Lendemor authored Dec 21, 2024
1 parent a2ec1bc commit 848b870
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 35 deletions.
22 changes: 12 additions & 10 deletions reflex/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -1356,20 +1356,22 @@ async def health() -> JSONResponse:
health_status = {"status": True}
status_code = 200

db_status, redis_status = await asyncio.gather(
get_db_status(), prerequisites.get_redis_status()
)
tasks = []

if prerequisites.check_db_used():
tasks.append(get_db_status())
if prerequisites.check_redis_used():
tasks.append(prerequisites.get_redis_status())

results = await asyncio.gather(*tasks)

health_status["db"] = db_status
for result in results:
health_status |= result

if redis_status is None:
if "redis" in health_status and health_status["redis"] is None:
health_status["redis"] = False
else:
health_status["redis"] = redis_status

if not health_status["db"] or (
not health_status["redis"] and redis_status is not None
):
if not all(health_status.values()):
health_status["status"] = False
status_code = 503

Expand Down
8 changes: 3 additions & 5 deletions reflex/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,15 +141,13 @@ def get_async_engine(url: str | None) -> sqlalchemy.ext.asyncio.AsyncEngine:
return _ASYNC_ENGINE[url]


async def get_db_status() -> bool:
async def get_db_status() -> dict[str, bool]:
"""Checks the status of the database connection.
Attempts to connect to the database and execute a simple query to verify connectivity.
Returns:
bool: The status of the database connection:
- True: The database is accessible.
- False: The database is not accessible.
The status of the database connection.
"""
status = True
try:
Expand All @@ -159,7 +157,7 @@ async def get_db_status() -> bool:
except sqlalchemy.exc.OperationalError:
status = False

return status
return {"db": status}


SQLModelOrSqlAlchemy = Union[
Expand Down
27 changes: 21 additions & 6 deletions reflex/utils/prerequisites.py
Original file line number Diff line number Diff line change
Expand Up @@ -372,16 +372,13 @@ def parse_redis_url() -> str | dict | None:
return config.redis_url


async def get_redis_status() -> bool | None:
async def get_redis_status() -> dict[str, bool | None]:
"""Checks the status of the Redis connection.
Attempts to connect to Redis and send a ping command to verify connectivity.
Returns:
bool or None: The status of the Redis connection:
- True: Redis is accessible and responding.
- False: Redis is not accessible due to a connection error.
- None: Redis not used i.e redis_url is not set in rxconfig.
The status of the Redis connection.
"""
try:
status = True
Expand All @@ -393,7 +390,7 @@ async def get_redis_status() -> bool | None:
except exceptions.RedisError:
status = False

return status
return {"redis": status}


def validate_app_name(app_name: str | None = None) -> str:
Expand Down Expand Up @@ -1177,6 +1174,24 @@ def initialize_frontend_dependencies():
initialize_web_directory()


def check_db_used() -> bool:
"""Check if the database is used.
Returns:
True if the database is used.
"""
return bool(get_config().db_url)


def check_redis_used() -> bool:
"""Check if Redis is used.
Returns:
True if Redis is used.
"""
return bool(get_config().redis_url)


def check_db_initialized() -> bool:
"""Check if the database migrations are initialized.
Expand Down
52 changes: 38 additions & 14 deletions tests/units/test_health_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@
"mock_redis_client, expected_status",
[
# Case 1: Redis client is available and responds to ping
(Mock(ping=lambda: None), True),
(Mock(ping=lambda: None), {"redis": True}),
# Case 2: Redis client raises RedisError
(Mock(ping=lambda: (_ for _ in ()).throw(RedisError)), False),
(Mock(ping=lambda: (_ for _ in ()).throw(RedisError)), {"redis": False}),
# Case 3: Redis client is not used
(None, None),
(None, {"redis": None}),
],
)
async def test_get_redis_status(mock_redis_client, expected_status, mocker):
Expand All @@ -41,12 +41,12 @@ async def test_get_redis_status(mock_redis_client, expected_status, mocker):
"mock_engine, execute_side_effect, expected_status",
[
# Case 1: Database is accessible
(MagicMock(), None, True),
(MagicMock(), None, {"db": True}),
# Case 2: Database connection error (OperationalError)
(
MagicMock(),
sqlalchemy.exc.OperationalError("error", "error", "error"),
False,
{"db": False},
),
],
)
Expand Down Expand Up @@ -74,25 +74,49 @@ async def test_get_db_status(mock_engine, execute_side_effect, expected_status,

@pytest.mark.asyncio
@pytest.mark.parametrize(
"db_status, redis_status, expected_status, expected_code",
"db_enabled, redis_enabled, db_status, redis_status, expected_status, expected_code",
[
# Case 1: Both services are connected
(True, True, {"status": True, "db": True, "redis": True}, 200),
(True, True, True, True, {"status": True, "db": True, "redis": True}, 200),
# Case 2: Database not connected, Redis connected
(False, True, {"status": False, "db": False, "redis": True}, 503),
(True, True, False, True, {"status": False, "db": False, "redis": True}, 503),
# Case 3: Database connected, Redis not connected
(True, False, {"status": False, "db": True, "redis": False}, 503),
(True, True, True, False, {"status": False, "db": True, "redis": False}, 503),
# Case 4: Both services not connected
(False, False, {"status": False, "db": False, "redis": False}, 503),
(True, True, False, False, {"status": False, "db": False, "redis": False}, 503),
# Case 5: Database Connected, Redis not used
(True, None, {"status": True, "db": True, "redis": False}, 200),
(True, False, True, None, {"status": True, "db": True}, 200),
# Case 6: Database not used, Redis Connected
(False, True, None, True, {"status": True, "redis": True}, 200),
# Case 7: Both services not used
(False, False, None, None, {"status": True}, 200),
],
)
async def test_health(db_status, redis_status, expected_status, expected_code, mocker):
async def test_health(
db_enabled,
redis_enabled,
db_status,
redis_status,
expected_status,
expected_code,
mocker,
):
# Mock get_db_status and get_redis_status
mocker.patch("reflex.app.get_db_status", return_value=db_status)
mocker.patch(
"reflex.utils.prerequisites.get_redis_status", return_value=redis_status
"reflex.utils.prerequisites.check_db_used",
return_value=db_enabled,
)
mocker.patch(
"reflex.utils.prerequisites.check_redis_used",
return_value=redis_enabled,
)
mocker.patch(
"reflex.app.get_db_status",
return_value={"db": db_status},
)
mocker.patch(
"reflex.utils.prerequisites.get_redis_status",
return_value={"redis": redis_status},
)

# Call the async health function
Expand Down

0 comments on commit 848b870

Please sign in to comment.