Add healthcheck integration in lifespan

Author: alukach

src/stac_auth_proxy/app.py

@@ -12,6 +12,7 @@
 
 from .config import Settings
 from .handlers import HealthzHandler, ReverseProxyHandler
+from .lifespan import LifespanManager, ServerHealthCheck
 from .middleware import (
     AddProcessTimeHeaderMiddleware,
     ApplyCql2FilterMiddleware,
@@ -27,8 +28,22 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     """FastAPI Application Factory."""
     settings = settings or Settings()
 
+    upstream_urls = [
+        settings.upstream_url,
+        settings.oidc_discovery_internal_url or settings.oidc_discovery_url,
+    ]
+    lifespan = LifespanManager(
+        on_startup=(
+            [ServerHealthCheck(url=url) for url in upstream_urls]
+            if settings.wait_for_upstream
+            else []
+        )
+    )
+
     app = FastAPI(
         openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema
+        lifespan=lifespan,
+    )
 
     # Add catchall proxy handler
     proxy_handler = ReverseProxyHandler(upstream=str(settings.upstream_url))
@@ -59,7 +74,6 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
         app.add_middleware(ApplyCql2FilterMiddleware)
         app.add_middleware(
             BuildCql2FilterMiddleware,
-            # collections_filter=settings.collections_filter,
             items_filter=settings.items_filter(),
         )
 

src/stac_auth_proxy/config.py

@@ -36,6 +36,8 @@ class Settings(BaseSettings):
     oidc_discovery_url: HttpUrl
     oidc_discovery_internal_url: Optional[HttpUrl] = None
 
+    wait_for_upstream: bool = True
+
     # Endpoints
     healthz_prefix: str = Field(pattern=_PREFIX_PATTERN, default="/healthz")
     openapi_spec_endpoint: Optional[str] = Field(pattern=_PREFIX_PATTERN, default=None)

src/stac_auth_proxy/lifespan/LifespanManager.py

@@ -0,0 +1,37 @@
+"""Lifespan manager for FastAPI applications."""
+
+import logging
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import AsyncGenerator, Awaitable, Callable, List
+
+from fastapi import FastAPI
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LifespanManager:
+    """Manager for FastAPI lifespan events."""
+
+    on_startup: List[Callable[[], Awaitable[None]]] = field(default_factory=list)
+    on_teardown: List[Callable[[], Awaitable[None]]] = field(default_factory=list)
+
+    @asynccontextmanager
+    async def __call__(self, app: FastAPI) -> AsyncGenerator[None, None]:
+        """FastAPI lifespan event handler."""
+        for i, task in enumerate(self.on_startup):
+            logger.debug(f"Executing startup task {i+1}/{len(self.on_startup)}")
+            await task()
+
+        logger.debug("All startup tasks completed successfully")
+
+        yield
+
+        # Execute teardown tasks
+        for i, task in enumerate(self.on_teardown):
+            try:
+                logger.debug(f"Executing teardown task {i+1}/{len(self.on_teardown)}")
+                await task()
+            except Exception as e:
+                logger.error(f"Teardown task failed: {e}")

src/stac_auth_proxy/lifespan/ServerHealthCheck.py

@@ -0,0 +1,57 @@
+"""Health check implementations for lifespan events."""
+
+import asyncio
+import logging
+from dataclasses import dataclass
+
+import httpx
+from pydantic import HttpUrl
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ServerHealthCheck:
+    """Health check for upstream API."""
+
+    url: str | HttpUrl
+    max_retries: int = 5
+    retry_delay: float = 0.25
+    retry_delay_max: float = 10.0
+    timeout: float = 5.0
+
+    def __post_init__(self):
+        """Convert url to string if it's a HttpUrl."""
+        if isinstance(self.url, HttpUrl):
+            self.url = str(self.url)
+
+    async def _check_health(self) -> bool:
+        """Check if upstream API is responding."""
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    self.url, timeout=self.timeout, follow_redirects=True
+                )
+                response.raise_for_status()
+                return True
+        except Exception as e:
+            logger.warning(f"Upstream health check for {self.url!r} failed: {e}")
+            return False
+
+    async def __call__(self) -> None:
+        """Wait for upstream API to become available."""
+        for attempt in range(self.max_retries):
+            if await self._check_health():
+                logger.info(f"Upstream API {self.url!r} is healthy")
+                return
+
+            retry_in = min(self.retry_delay * (2**attempt), self.retry_delay_max)
+            logger.warning(
+                f"Upstream API {self.url!r} not healthy, retrying in {retry_in:.1f}s "
+                f"(attempt {attempt + 1}/{self.max_retries})"
+            )
+            await asyncio.sleep(retry_in)
+
+        raise RuntimeError(
+            f"Upstream API {self.url!r} failed to respond after {self.max_retries} attempts"
+        )

src/stac_auth_proxy/lifespan/__init__.py

@@ -0,0 +1,9 @@
+"""Lifespan event handlers for the STAC Auth Proxy."""
+
+from .LifespanManager import LifespanManager
+from .ServerHealthCheck import ServerHealthCheck
+
+__all__ = [
+    "ServerHealthCheck",
+    "LifespanManager",
+]