fix: handle Starlette routes in FastAPI middleware

martin1keogh · martin1keogh · commit 75d6a5a797cf · 2022-06-21T14:12:18.000+02:00
This commit changes the FastAPI middleware so that the `"route"` info is now optional. If it is not present, the SQL comment will not have the controller/route info. This info is always present in FastAPI routes. However, Starlette-based routes (FastAPI can mount Starlette apps as sub-applications) do not and would crash when the middleware tried to extract the information. Fixes #133
diff --git a/python/sqlcommenter-python/google/cloud/sqlcommenter/fastapi.py b/python/sqlcommenter-python/google/cloud/sqlcommenter/fastapi.py
@@ -95,5 +95,5 @@ def _get_fastapi_route(fastapi_app: FastAPI, scope) -> Optional[Route]:
         # and return the route name if found.
         match, child_scope = route.matches(scope)
         if match == Match.FULL:
-            return child_scope["route"]
+            return child_scope.get("route")
     return None
diff --git a/python/sqlcommenter-python/tests/fastapi/app.py b/python/sqlcommenter-python/tests/fastapi/app.py
@@ -2,7 +2,9 @@
 
 from fastapi import FastAPI, status
 from fastapi.responses import JSONResponse
+from starlette.applications import Starlette
 from starlette.exceptions import HTTPException as StarletteHTTPException
+from starlette.routing import Route
 
 from google.cloud.sqlcommenter.fastapi import SQLCommenterMiddleware, get_fastapi_info
 
@@ -27,3 +29,15 @@ async def custom_http_exception_handler(request, exc):
         status_code=status.HTTP_404_NOT_FOUND,
         content=get_fastapi_info(),
     )
+
+
+def starlette_endpoint(_):
+    return JSONResponse({"from": "starlette"})
+
+
+starlette_subapi = Starlette(routes=[
+    Route("/", starlette_endpoint),
+])
+
+
+app.mount("/starlette", starlette_subapi)
diff --git a/python/sqlcommenter-python/tests/fastapi/tests.py b/python/sqlcommenter-python/tests/fastapi/tests.py
@@ -53,3 +53,13 @@ def test_get_fastapi_info_in_404_error_context(client):
 
 def test_get_fastapi_info_outside_request_context(client):
     assert get_fastapi_info() == {}
+
+
+def test_get_openapi_does_not_throw_an_error(client):
+    resp = client.get(app.docs_url)
+    assert resp.status_code == 200
+
+
+def test_get_starlette_endpoints_does_not_throw_an_error(client):
+    resp = client.get("/starlette")
+    assert resp.status_code == 200
