fix health check and skip not needed tasks (#4563)
This commit is contained in:
parent
a2ec1bc1d8
commit
848b87070c
@ -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 = []
|
||||
|
||||
health_status["db"] = db_status
|
||||
if prerequisites.check_db_used():
|
||||
tasks.append(get_db_status())
|
||||
if prerequisites.check_redis_used():
|
||||
tasks.append(prerequisites.get_redis_status())
|
||||
|
||||
if redis_status is None:
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
for result in results:
|
||||
health_status |= result
|
||||
|
||||
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
|
||||
|
||||
|
@ -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:
|
||||
@ -159,7 +157,7 @@ async def get_db_status() -> bool:
|
||||
except sqlalchemy.exc.OperationalError:
|
||||
status = False
|
||||
|
||||
return status
|
||||
return {"db": status}
|
||||
|
||||
|
||||
SQLModelOrSqlAlchemy = Union[
|
||||
|
@ -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
|
||||
@ -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:
|
||||
@ -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.
|
||||
|
||||
|
@ -15,11 +15,11 @@ from reflex.utils.prerequisites import get_redis_status
|
||||
"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):
|
||||
@ -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},
|
||||
),
|
||||
],
|
||||
)
|
||||
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user