Merge pull request #1431 from PrefectHQ/can-connect-fixup

Avoid hiding the connection error when the api healthcheck fails

src/prefect/client.py

@@ -245,13 +245,19 @@ def api_url(self) -> httpx.URL:
 
     # API methods ----------------------------------------------------------------------
 
-    async def api_healthcheck(self) -> bool:
+    async def api_healthcheck(self) -> Optional[Exception]:
+        """
+        Attempts to connect to the API and returns the encountered exception if not
+        successful.
+
+        If successful, returns `None`.
+        """
         try:
             with anyio.fail_after(10):
                 await self._client.get("/health")
-                return True
-        except:
-            return False
+                return None
+        except Exception as exc:
+            return exc
 
     async def hello(self) -> httpx.Response:
         """

src/prefect/engine.py

@@ -140,11 +140,11 @@ async def create_then_begin_flow_run(
 
     Creates the flow run in the backend, then enters the main flow run engine.
     """
-    can_connect = await client.api_healthcheck()
-    if not can_connect:
+    connect_error = await client.api_healthcheck()
+    if connect_error:
         raise RuntimeError(
             f"Cannot create flow run. Failed to reach API at {client.api_url}."
-        )
+        ) from connect_error
 
     state = Pending()
     if flow.should_validate_parameters:
@@ -652,12 +652,12 @@ async def begin_task_run(
             # Otherwise, retrieve a new client
             client = await stack.enter_async_context(get_client())
 
-        can_connect = await client.api_healthcheck()
-        if not can_connect:
+        connect_error = await client.api_healthcheck()
+        if connect_error:
             raise RuntimeError(
                 f"Cannot orchestrate task run '{task_run.id}'. "
                 f"Failed to connect to API at {client.api_url}."
-            )
+            ) from connect_error
 
         return await orchestrate_task_run(
             task=task,

tests/test_client.py

@@ -23,7 +23,7 @@
 from prefect.orion.schemas.schedules import IntervalSchedule
 from prefect.orion.schemas.states import Pending, Running, Scheduled, StateType
 from prefect.tasks import task
-from prefect.utilities.testing import AsyncMock, temporary_settings
+from prefect.utilities.testing import AsyncMock, exceptions_equal, temporary_settings
 
 
 class TestGetClient:
@@ -214,7 +214,14 @@ async def test_hello(orion_client):
 
 
 async def test_healthcheck(orion_client):
-    assert await orion_client.api_healthcheck()
+    assert await orion_client.api_healthcheck() is None
+
+
+async def test_healthcheck_failure(orion_client, monkeypatch):
+    monkeypatch.setattr(
+        orion_client._client, "get", AsyncMock(side_effect=ValueError("test"))
+    )
+    assert exceptions_equal(await orion_client.api_healthcheck(), ValueError("test"))
 
 
 async def test_create_then_read_flow(orion_client):