From f8a111fce780996a4c3a94b5637cdbe0142d1961 Mon Sep 17 00:00:00 2001
From: Iain-S <25081046+Iain-S@users.noreply.github.com>
Date: Wed, 3 Jul 2024 13:51:26 +0100
Subject: [PATCH] Switch to fastapi lifespan for startup/shutdown handling

---
 rctab/main.py | 51 ++++++++++++++++++++++++++-------------------------
 1 file changed, 26 insertions(+), 25 deletions(-)

diff --git a/rctab/main.py b/rctab/main.py
index 5d22dea..28eeb56 100644
--- a/rctab/main.py
+++ b/rctab/main.py
@@ -1,8 +1,9 @@
 """The entrypoint of the FastAPI application."""
 
 import logging
+from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, Callable, Dict, Final
+from typing import Any, AsyncIterator, Callable, Dict, Final
 
 import fastapimsal
 import secure
@@ -34,6 +35,29 @@
 
 templates = Jinja2Templates(directory=Path("rctab/templates"))
 
+
+@asynccontextmanager
+async def lifespan(_: FastAPI) -> AsyncIterator[None]:
+    """Handle setup and teardown."""
+    await database.connect()
+    settings = get_settings()
+    logging.basicConfig(level=settings.log_level)
+    set_log_handler()
+    if not settings.ignore_whitelist:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "Starting server with subscription whitelist: %s", settings.whitelist
+        )
+
+    yield
+
+    logger = logging.getLogger(__name__)
+    logger.warning("Shutting down server...")
+
+    logger.info("Disconnecting from database")
+    await database.disconnect()
+
+
 app = FastAPI(
     title="RCTab API",
     description="API for RCTab",
@@ -41,6 +65,7 @@
     docs_url=None,
     redoc_url=None,
     openapi_url=None,
+    lifespan=lifespan,
 )
 
 server = secure.Server().set("Secure")
@@ -76,30 +101,6 @@ async def set_secure_headers(request: Any, call_next: Callable[[Any], Any]) -> A
 )
 
 
-@app.on_event("startup")
-async def startup() -> None:
-    """Start the server up."""
-    await database.connect()
-    settings = get_settings()
-    logging.basicConfig(level=settings.log_level)
-    set_log_handler()
-    if not settings.ignore_whitelist:
-        logger = logging.getLogger(__name__)
-        logger.warning(
-            "Starting server with subscription whitelist: %s", settings.whitelist
-        )
-
-
-@app.on_event("shutdown")
-async def shutdown() -> None:
-    """Shut the server down."""
-    logger = logging.getLogger(__name__)
-    logger.warning("Shutting down server...")
-
-    logger.info("Disconnecting from database")
-    await database.disconnect()
-
-
 @app.exception_handler(UniqueViolationError)
 async def unicorn_exception_handler(
     _: Request, exc: UniqueViolationError