Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix health check and skip not needed tasks #4563

Merged
merged 2 commits into from
Dec 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading