better handle sync/async health check functions (#864)

Author: vincentsarago

CHANGES.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed 
+
+- better support async/sync health check functions
+
 ### Removed
 
 - support for python 3.9 and 3.10

stac_fastapi/api/stac_fastapi/api/app.py

@@ -1,14 +1,14 @@
 """Fastapi app creation."""
 
 
+import inspect
 from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import attr
 from brotli_asgi import BrotliMiddleware
 from fastapi import APIRouter, FastAPI
 from fastapi.params import Depends
 from stac_pydantic import api
-from stac_pydantic.api.collections import Collections
 from stac_pydantic.shared import MimeTypes
 from stac_pydantic.version import STAC_VERSION
 from starlette.middleware import Middleware
@@ -32,6 +32,7 @@
     add_direct_response,
     add_route_dependencies,
     create_async_endpoint,
+    sync_to_async,
 )
 from stac_fastapi.types.config import ApiSettings, Settings
 from stac_fastapi.types.core import AsyncBaseCoreClient, BaseCoreClient
@@ -277,14 +278,14 @@ def register_get_collections(self) -> None:
             name="Get Collections",
             path="/collections",
             response_model=(
-                Collections if self.settings.enable_response_models else None
+                api.Collections if self.settings.enable_response_models else None
             ),
             responses={
                 200: {
                     "content": {
                         MimeTypes.json.value: {},
                     },
-                    "model": Collections,
+                    "model": api.Collections,
                 },
             },
             response_class=self.response_class,
@@ -411,7 +412,9 @@ async def ping():
             },
             response_class=self.response_class,
             methods=["GET"],
-            endpoint=self.health_check,
+            endpoint=self.health_check
+            if inspect.iscoroutinefunction(self.health_check)
+            else sync_to_async(self.health_check),
         )
         self.app.include_router(mgmt_router, tags=["Liveliness/Readiness"])
 

stac_fastapi/api/tests/test_app.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import List, Optional, Union
 
 import attr
@@ -578,3 +579,34 @@ def health_check(request: Request):
                 "version": "0.1.0",
             },
         }
+
+    async def async_health_check(request: Request):
+        await asyncio.sleep(1)
+        return {
+            "status": "UP",
+            "database": {
+                "status": "UP",
+                "version": "0.1.0",
+            },
+        }
+
+    test_app = app.StacApi(
+        settings=ApiSettings(),
+        client=AsyncTestCoreClient(),
+        health_check=async_health_check,
+    )
+
+    with TestClient(test_app.app) as client:
+        resp = client.get("/_mgmt/ping")
+        assert resp.status_code == 200
+        assert resp.json() == {"message": "PONG"}
+
+        resp = client.get("/_mgmt/health")
+        assert resp.status_code == 200
+        assert resp.json() == {
+            "status": "UP",
+            "database": {
+                "status": "UP",
+                "version": "0.1.0",
+            },
+        }