Rework lifespan

Author: alukach

src/stac_auth_proxy/__init__.py

@@ -8,13 +8,9 @@
 
 from .app import configure_app, create_app
 from .config import Settings
-from .lifespan import check_conformance, check_server_health, lifespan
 
 __all__ = [
     "create_app",
     "configure_app",
-    "lifespan",
-    "check_conformance",
-    "check_server_health",
     "Settings",
 ]

src/stac_auth_proxy/app.py

@@ -13,7 +13,7 @@
 
 from .config import Settings
 from .handlers import HealthzHandler, ReverseProxyHandler, SwaggerUI
-from .lifespan import lifespan
+from .lifespan import build_lifespan
 from .middleware import (
     AddProcessTimeHeaderMiddleware,
     AuthenticationExtensionMiddleware,
@@ -31,11 +31,11 @@
 
 
 def configure_app(app: FastAPI, settings: Optional[Settings] = None) -> FastAPI:
-    """Apply routes and middleware to an existing FastAPI app."""
+    """Apply routes and middleware to a FastAPI app."""
     settings = settings or Settings()
 
     #
-    # Handlers (place catch-all proxy handler last)
+    # Route Handlers
     #
 
     # If we have customized Swagger UI Init settings (e.g. a provided client_id)
@@ -143,7 +143,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
 
     app = FastAPI(
         openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema
-        lifespan=lifespan(settings=settings),
+        lifespan=build_lifespan(settings=settings),
         root_path=settings.root_path,
     )
     if app.root_path:

src/stac_auth_proxy/lifespan.py

@@ -1,18 +1,116 @@
 """Reusable lifespan handler for FastAPI applications."""
 
+import asyncio
 import logging
+import re
 from contextlib import asynccontextmanager
 from typing import Any
 
+import httpx
 from fastapi import FastAPI
+from pydantic import HttpUrl
+from starlette.middleware import Middleware
 
 from .config import Settings
-from .utils.lifespan import check_conformance, check_server_health
 
 logger = logging.getLogger(__name__)
+__all__ = ["build_lifespan", "check_conformance", "check_server_health"]
+
+
+async def check_server_healths(*urls: str | HttpUrl) -> None:
+    """Wait for upstream APIs to become available."""
+    logger.info("Running upstream server health checks...")
+    for url in urls:
+        await check_server_health(url)
+    logger.info(
+        "Upstream servers are healthy:\n%s",
+        "\n".join([f" - {url}" for url in urls]),
+    )
+
+
+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(
+        base_url=url, timeout=timeout, follow_redirects=True
+    ) as client:
+        for attempt in range(max_retries):
+            try:
+                response = await client.get("/")
+                response.raise_for_status()
+                logger.info(f"Upstream API {url!r} is healthy")
+                return
+            except httpx.ConnectError 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"
+    )
+
+
+async def check_conformance(
+    middleware_classes: list[Middleware],
+    api_url: str,
+    attr_name: str = "__required_conformances__",
+    endpoint: str = "/conformance",
+):
+    """Check if the upstream API supports a given conformance class."""
+    required_conformances: dict[str, list[str]] = {}
+    for middleware in middleware_classes:
+
+        for conformance in getattr(middleware.cls, attr_name, []):
+            required_conformances.setdefault(conformance, []).append(
+                middleware.cls.__name__
+            )
+
+    async with httpx.AsyncClient(base_url=api_url) as client:
+        response = await client.get(endpoint)
+        response.raise_for_status()
+        api_conforms_to = response.json().get("conformsTo", [])
+
+    missing = [
+        req_conformance
+        for req_conformance in required_conformances.keys()
+        if not any(
+            re.match(req_conformance, conformance) for conformance in api_conforms_to
+        )
+    ]
+
+    def conformance_str(conformance: str) -> str:
+        return f" - {conformance} [{','.join(required_conformances[conformance])}]"
+
+    if missing:
+        missing_str = [conformance_str(c) for c in missing]
+        raise RuntimeError(
+            "\n".join(
+                [
+                    "Upstream catalog is missing the following conformance classes:",
+                    *missing_str,
+                ]
+            )
+        )
+    logger.info(
+        "Upstream catalog conforms to the following required conformance classes: \n%s",
+        "\n".join([conformance_str(c) for c in required_conformances]),
+    )
 
 
-def lifespan(settings: Settings | None = None, **settings_kwargs: Any):
+def build_lifespan(settings: Settings | None = None, **settings_kwargs: Any):
     """
     Create a lifespan handler that runs startup checks.
 
@@ -34,18 +132,13 @@ def lifespan(settings: Settings | None = None, **settings_kwargs: Any):
         settings = Settings(**settings_kwargs)
 
     @asynccontextmanager
-    async def _lifespan(app: FastAPI):
+    async def lifespan(app: "FastAPI"):
         assert settings is not None  # Required for type checking
 
         # Wait for upstream servers to become available
         if settings.wait_for_upstream:
-            logger.info("Running upstream server health checks...")
-            urls = [settings.upstream_url, settings.oidc_discovery_internal_url]
-            for url in urls:
-                await check_server_health(url=url)
-            logger.info(
-                "Upstream servers are healthy:\n%s",
-                "\n".join([f" - {url}" for url in urls]),
+            await check_server_healths(
+                settings.upstream_url, settings.oidc_discovery_internal_url
             )
 
         # Log all middleware connected to the app
@@ -59,7 +152,4 @@ async def _lifespan(app: FastAPI):
 
         yield
 
-    return _lifespan
-
-
-__all__ = ["lifespan", "check_conformance", "check_server_health"]
+    return lifespan