refactor: simplify lifespan

Author: alukach

src/stac_auth_proxy/app.py

@@ -6,21 +6,22 @@
 """
 
 import logging
+from contextlib import asynccontextmanager
 from typing import Optional
 
 from fastapi import FastAPI
 from starlette_cramjam.middleware import CompressionMiddleware
 
 from .config import Settings
 from .handlers import HealthzHandler, ReverseProxyHandler
-from .lifespan import LifespanManager, ServerHealthCheck
 from .middleware import (
     AddProcessTimeHeaderMiddleware,
     ApplyCql2FilterMiddleware,
     BuildCql2FilterMiddleware,
     EnforceAuthMiddleware,
     OpenApiMiddleware,
 )
+from .utils.lifespan import check_server_health
 
 logger = logging.getLogger(__name__)
 
@@ -32,14 +33,17 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     #
     # Application
     #
-    upstream_urls = (
-        [settings.upstream_url, settings.oidc_discovery_internal_url]
-        if settings.wait_for_upstream
-        else []
-    )
-    lifespan = LifespanManager(
-        on_startup=([ServerHealthCheck(url=url) for url in upstream_urls])
-    )
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        assert settings
+
+        # Wait for upstream servers to become available
+        if settings.wait_for_upstream:
+            for url in [settings.upstream_url, settings.oidc_discovery_internal_url]:
+                await check_server_health(url=url)
+
+        yield
 
     app = FastAPI(
         openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema

src/stac_auth_proxy/lifespan/LifespanManager.py

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