From 848b87070c1270af7c12d6b8d300f0dae0d71462 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Brand=C3=A9ho?= Date: Fri, 20 Dec 2024 17:08:10 -0800 Subject: [PATCH] fix health check and skip not needed tasks (#4563) --- reflex/app.py | 22 ++++++------ reflex/model.py | 8 ++--- reflex/utils/prerequisites.py | 27 +++++++++++---- tests/units/test_health_endpoint.py | 52 +++++++++++++++++++++-------- 4 files changed, 74 insertions(+), 35 deletions(-) diff --git a/reflex/app.py b/reflex/app.py index 935fe7900c..032b198f6d 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -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 diff --git a/reflex/model.py b/reflex/model.py index cb8bed29bb..de63589fc5 100644 --- a/reflex/model.py +++ b/reflex/model.py @@ -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[ diff --git a/reflex/utils/prerequisites.py b/reflex/utils/prerequisites.py index e5ff3185e5..797d28701b 100644 --- a/reflex/utils/prerequisites.py +++ b/reflex/utils/prerequisites.py @@ -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. diff --git a/tests/units/test_health_endpoint.py b/tests/units/test_health_endpoint.py index fe350266ff..6d12d79d6a 100644 --- a/tests/units/test_health_endpoint.py +++ b/tests/units/test_health_endpoint.py @@ -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): @@ -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