fix(server): Enable data platform connections

Author: devsjc

src/quartz_api/cmd/main.py

@@ -27,11 +27,14 @@
 ```
 """
 
+import functools
 import importlib
 import importlib.metadata
 import logging
 import pathlib
 import sys
+from collections.abc import Generator
+from contextlib import asynccontextmanager
 from typing import Any
 
 import uvicorn
@@ -41,7 +44,7 @@
 from fastapi.openapi.utils import get_openapi
 from grpclib.client import Channel
 from pydantic import BaseModel
-from pyhocon import ConfigFactory
+from pyhocon import ConfigFactory, ConfigTree
 from starlette.responses import FileResponse
 from starlette.staticfiles import StaticFiles
 
@@ -59,6 +62,7 @@ class GetHealthResponse(BaseModel):
 
     status: int
 
+
 def _custom_openapi(server: FastAPI) -> dict[str, Any]:
     """Customize the OpenAPI schema for ReDoc."""
     if server.openapi_schema:
@@ -87,20 +91,51 @@ def _custom_openapi(server: FastAPI) -> dict[str, Any]:
     return openapi_schema
 
 
-def run() -> None:
-    """Run the API using a uvicorn server."""
-    # Get the application configuration from the environment
-    conf = ConfigFactory.parse_file((pathlib.Path(__file__).parent / "server.conf").as_posix())
+@asynccontextmanager
+async def _lifespan(server: FastAPI, conf: ConfigTree) -> Generator[None]:
+    """Configure FastAPI app instance with startup and shutdown events."""
+    db_instance: models.DatabaseInterface | None = None
+    grpc_channel: Channel | None = None
+
+    match conf.get_string("backend.source"):
+        case "quartzdb":
+            db_instance = QuartzClient(
+                database_url=conf.get_string("backend.quartzdb.database_url"),
+            )
+        case "dummydb":
+            db_instance = DummyClient()
+            log.warning("disabled backend. NOT recommended for production")
+        case "dataplatform":
+            grpc_channel = Channel(
+                host=conf.get_string("backend.dataplatform.host"),
+                port=conf.get_int("backend.dataplatform.port"),
+            )
+            client = dp.DataPlatformDataServiceStub(channel=grpc_channel)
+            db_instance = DataPlatformClient.from_dp(dp_client=client)
+        case _ as backend_type:
+            raise ValueError(f"Unknown backend: {backend_type}")
 
-    # Create the FastAPI server
+    server.dependency_overrides[models.get_db_client] = lambda: db_instance
+
+    yield
+
+    if grpc_channel:
+        grpc_channel.close()
+
+
+def _create_server(conf: ConfigTree) -> FastAPI:
+    """Configure FastAPI app instance with routes, dependencies, and middleware."""
     server = FastAPI(
         version=importlib.metadata.version("quartz_api"),
+        lifespan=functools.partial(_lifespan, conf=conf),
         title="Quartz API",
         description=__doc__,
-        openapi_tags=[{
-            "name": "API Information",
-            "description": "Routes providing information about the API.",
-        }],
+        openapi_tags=[
+            {
+                "name": "API Information",
+                "description": "Routes providing information about the API.",
+            },
+        ],
         docs_url="/swagger",
         redoc_url=None,
     )
@@ -126,15 +161,34 @@ def redoc_html() -> FileResponse:
     # Setup sentry, if configured
     if conf.get_string("sentry.dsn") != "":
         import sentry_sdk
+
         sentry_sdk.init(
             dsn=conf.get_string("sentry.dsn"),
             environment=conf.get_string("sentry.environment"),
             traces_sample_rate=1,
         )
 
-        sentry_sdk.set_tag("app_name", "quartz_api")
+        sentry_sdk.set_tag("server_name", "quartz_api")
         sentry_sdk.set_tag("version", importlib.metadata.version("quartz_api"))
 
+    # Add routers to the server according to configuration
+    for r in conf.get_string("api.routers").split(","):
+        try:
+            mod = importlib.import_module(service.__name__ + f".{r}")
+            server.include_router(mod.router)
+        except ModuleNotFoundError as e:
+            raise OSError(f"No such router router '{r}'") from e
+        server.openapi_tags = [
+            {
+                "name": mod.__name__.split(".")[-1].capitalize(),
+                "description": mod.__doc__,
+            },
+            *server.openapi_tags,
+        ]
+
+    # Customize the OpenAPI schema
+    server.openapi = lambda: _custom_openapi(server)
+
     # Override dependencies according to configuration
     match (conf.get_string("auth0.domain"), conf.get_string("auth0.audience")):
         case (_, "") | ("", _):
@@ -149,32 +203,6 @@ def redoc_html() -> FileResponse:
         case _:
             raise ValueError("Invalid Auth0 configuration")
 
-    db_instance: models.DatabaseInterface
-    match conf.get_string("backend.source"):
-        case "quartzdb":
-            db_instance = QuartzClient(
-                database_url=conf.get_string("backend.quartzdb.database_url"),
-            )
-        case "dummydb":
-            db_instance = DummyClient()
-            log.warning("disabled backend. NOT recommended for production")
-        case "dataplatform":
-
-            channel = Channel(
-                host=conf.get_string("backend.dataplatform.host"),
-                port=conf.get_int("backend.dataplatform.port"),
-            )
-            client = dp.DataPlatformDataServiceStub(channel=channel)
-            db_instance = DataPlatformClient.from_dp(dp_client=client)
-        case _:
-            raise ValueError(
-                "Unknown backend. "
-                f"Expected one of {list(conf.get('backend').keys())}",
-            )
-
-    server.dependency_overrides[models.get_db_client] = lambda: db_instance
-
-    # Add IANA timezone dependency
     timezone: str = conf.get_string("api.timezone")
     server.dependency_overrides[models.get_timezone] = lambda: timezone
 
@@ -186,34 +214,27 @@ def redoc_html() -> FileResponse:
         allow_methods=["*"],
         allow_headers=["*"],
     )
-    server.add_middleware(
-        audit.RequestLoggerMiddleware,
-        db_client=db_instance,
-    )
+    server.add_middleware(audit.RequestLoggerMiddleware)
 
-    # Add routers to the server according to configuration
-    for r in conf.get_string("api.routers").split(","):
-        try:
-            mod = importlib.import_module(service.__name__ + f".{r}")
-            server.include_router(mod.router)
-        except ModuleNotFoundError as e:
-            raise OSError(f"No such router router '{r}'") from e
-        server.openapi_tags = [{
-            "name": mod.__name__.split(".")[-1].capitalize(),
-            "description": mod.__doc__,
-        }, *server.openapi_tags]
+    return server
 
-    # Customize the OpenAPI schema
-    server.openapi = lambda: _custom_openapi(server)
+
+def run() -> None:
+    """Run the API using a uvicorn server."""
+    # Get the application configuration from the environment
+    conf = ConfigFactory.parse_file((pathlib.Path(__file__).parent / "server.conf").as_posix())
+
+    server = _create_server(conf=conf)
 
     # Run the server with uvicorn
     uvicorn.run(
         server,
-        host="0.0.0.0", # noqa: S104
+        host="0.0.0.0",  # noqa: S104
         port=conf.get_int("api.port"),
         log_level=conf.get_string("api.loglevel"),
     )
 
 
 if __name__ == "__main__":
     run()
+

src/quartz_api/internal/middleware/audit.py

@@ -12,10 +12,9 @@
 class RequestLoggerMiddleware(BaseHTTPMiddleware):
     """Middleware to log API requests to the database."""
 
-    def __init__(self, server: FastAPI, db_client: models.DatabaseInterface) -> None:
+    def __init__(self, server: FastAPI) -> None:
         """Initialize the middleware with the FastAPI server and database client."""
         super().__init__(server)
-        self.db_client = db_client
 
     async def dispatch(
         self,
@@ -38,6 +37,9 @@ async def dispatch(
              url += f"?{request.url.query}"
 
         try:
+            db_client: models.DatabaseInterface = getattr(request.app.state, "db_instance", None)
+            if db_client is None:
+                raise RuntimeError("Database client not found in app state.")
             await self.db_client.save_api_call_to_db(url=url, authdata=auth)
         except Exception as e:
             logging.error(f"Failed to log request to DB: {e}")

src/quartz_api/internal/service/substations/router.py

@@ -11,7 +11,6 @@
 
 router = APIRouter(tags=[pathlib.Path(__file__).parent.stem.capitalize()])
 
-
 @router.get(
     "/substations",
     status_code=status.HTTP_200_OK,