From 484494b1947cba04d99c9068c12e05a39f555f7f Mon Sep 17 00:00:00 2001 From: Martin Keogh Date: Tue, 21 Jun 2022 14:04:31 +0200 Subject: [PATCH] fix: handle Starlette routes in FastAPI middleware 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 --- .../google/cloud/sqlcommenter/fastapi.py | 2 +- python/sqlcommenter-python/tests/fastapi/app.py | 14 ++++++++++++++ python/sqlcommenter-python/tests/fastapi/tests.py | 10 ++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/python/sqlcommenter-python/google/cloud/sqlcommenter/fastapi.py b/python/sqlcommenter-python/google/cloud/sqlcommenter/fastapi.py index 3013492a..0fcc5f36 100644 --- 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 index 24dcc8ce..6bb88625 100644 --- a/python/sqlcommenter-python/tests/fastapi/app.py +++ b/python/sqlcommenter-python/tests/fastapi/app.py @@ -5,7 +5,9 @@ from google.cloud.sqlcommenter.fastapi import ( SQLCommenterMiddleware, get_fastapi_info, ) +from starlette.applications import Starlette from starlette.exceptions import HTTPException as StarletteHTTPException +from starlette.routing import Route app = FastAPI(title="SQLCommenter") @@ -28,3 +30,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 index d2edb9f5..2fb84945 100644 --- a/python/sqlcommenter-python/tests/fastapi/tests.py +++ b/python/sqlcommenter-python/tests/fastapi/tests.py @@ -52,3 +52,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