Skip to content

Commit e1d9a49

Browse files
author
Jose Luis Navarro
authored
Add FastAPI support with an ASGI Middleware (#105)
* Add FastAPI support with an ASGI Middleware * Fix PR comments
1 parent 181a4c6 commit e1d9a49

8 files changed

Lines changed: 241 additions & 1 deletion

File tree

‎python/sqlcommenter-python/README.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,20 @@ traceparent='00-5bd66ef5095369c7b0d1f8f4bd33716a-c532cb4098ac3dd2-01',
8585
tracestate='congo%%3Dt61rcWkgMzE%%2Crojo%%3D00f067aa0ba902b7'*/
8686
```
8787

88+
#### FastAPI
89+
90+
If you are using SQLAlchemy with FastAPI, add the middleware to get: framework, app_name, controller and route.
91+
92+
```python
93+
from fastapi import FastAPI
94+
from google.cloud.sqlcommenter.fastapi import SQLCommenterMiddleware
95+
96+
app = FastAPI()
97+
98+
app.add_middleware(SQLCommenterMiddleware)
99+
```
100+
101+
88102
### Psycopg2
89103

90104
Use the provided cursor factory to generate database cursors. All queries executed with such cursors will have the SQL comment prepended to them.

‎python/sqlcommenter-python/google/cloud/sqlcommenter/fastapi.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#!/usr/bin/python
2+
#
3+
# Copyright 2022 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
from __future__ import absolute_import
18+
19+
try:
20+
from typing import Optional
21+
from asgiref.compatibility import guarantee_single_callable
22+
from contextvars import ContextVar
23+
import fastapi
24+
from fastapi import FastAPI
25+
from starlette.routing import Match, Route
26+
except ImportError:
27+
fastapi = None
28+
29+
context = ContextVar("context", default={})
30+
31+
32+
def get_fastapi_info():
33+
"""
34+
Get available info from the current FastAPI request, if we're in a
35+
FastAPI request-response cycle. Else, return an empty dict.
36+
"""
37+
info = {}
38+
if fastapi and context:
39+
info = context.get()
40+
return info
41+
42+
43+
class SQLCommenterMiddleware:
44+
"""The ASGI application middleware.
45+
This class is an ASGI middleware that augment SQL statements before execution,
46+
with comments containing information about the code that caused its execution.
47+
48+
Args:
49+
app: The ASGI application callable to forward requests to.
50+
"""
51+
52+
def __init__(self, app):
53+
self.app = guarantee_single_callable(app)
54+
55+
async def __call__(self, scope, receive, send):
56+
"""The ASGI application
57+
Args:
58+
scope: An ASGI environment.
59+
receive: An awaitable callable yielding dictionaries
60+
send: An awaitable callable taking a single dictionary as argument.
61+
"""
62+
if scope["type"] not in ("http", "websocket"):
63+
return await self.app(scope, receive, send)
64+
65+
if not isinstance(scope["app"], FastAPI):
66+
return await self.app(scope, receive, send)
67+
68+
fastapi_app = scope["app"]
69+
info = _get_fastapi_info(fastapi_app, scope)
70+
token = context.set(info)
71+
72+
try:
73+
await self.app(scope, receive, send)
74+
finally:
75+
context.reset(token)
76+
77+
78+
def _get_fastapi_info(fastapi_app: FastAPI, scope) -> dict:
79+
info = {
80+
"framework": 'fastapi:%s' % fastapi.__version__,
81+
"app_name": fastapi_app.title,
82+
}
83+
84+
route = _get_fastapi_route(fastapi_app, scope)
85+
if route:
86+
info["controller"] = route.name
87+
info["route"] = route.path
88+
89+
return info
90+
91+
92+
def _get_fastapi_route(fastapi_app: FastAPI, scope) -> Optional[Route]:
93+
for route in fastapi_app.router.routes:
94+
# Determine if any route matches the incoming scope,
95+
# and return the route name if found.
96+
match, child_scope = route.matches(scope)
97+
if match == Match.FULL:
98+
return child_scope["route"]
99+
return None

‎python/sqlcommenter-python/google/cloud/sqlcommenter/sqlalchemy/executor.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import sqlalchemy
2424
from google.cloud.sqlcommenter import generate_sql_comment
25+
from google.cloud.sqlcommenter.fastapi import get_fastapi_info
2526
from google.cloud.sqlcommenter.flask import get_flask_info
2627
from google.cloud.sqlcommenter.opencensus import get_opencensus_values
2728
from google.cloud.sqlcommenter.opentelemetry import get_opentelemetry_values
@@ -42,6 +43,12 @@ def BeforeExecuteFactory(
4243
'db_framework': with_db_framework,
4344
}
4445

46+
def get_framework_info():
47+
info = get_flask_info()
48+
if not info:
49+
info = get_fastapi_info()
50+
return info
51+
4552
def before_cursor_execute(conn, cursor, sql, parameters, context, executemany):
4653
data = dict(
4754
# TODO: Figure out how to retrieve the exact driver version.
@@ -54,7 +61,7 @@ def before_cursor_execute(conn, cursor, sql, parameters, context, executemany):
5461
# folks using it in a web framework such as flask will
5562
# use it in unison with flask but initialize the parts disjointly,
5663
# unlike Django which uses ORMs directly as part of the framework.
57-
data.update(get_flask_info())
64+
data.update(get_framework_info())
5865

5966
# Filter down to just the requested attributes.
6067
data = {k: v for k, v in data.items() if attributes.get(k)}

‎python/sqlcommenter-python/setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ def read_file(filename):
3737
extras_require={
3838
'django': ['django >= 1.11'],
3939
'flask': ['flask'],
40+
'fastapi': ['fastapi'],
4041
'psycopg2': ['psycopg2'],
4142
'sqlalchemy': ['sqlalchemy'],
4243
'opencensus': ['opencensus'],

‎python/sqlcommenter-python/tests/fastapi/__init__.py‎

Whitespace-only changes.

‎python/sqlcommenter-python/tests/fastapi/app.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from typing import Optional
2+
3+
from fastapi import FastAPI, status
4+
from fastapi.responses import JSONResponse
5+
from starlette.exceptions import HTTPException as StarletteHTTPException
6+
7+
from google.cloud.sqlcommenter.fastapi import SQLCommenterMiddleware, get_fastapi_info
8+
9+
app = FastAPI(title="SQLCommenter")
10+
11+
app.add_middleware(SQLCommenterMiddleware)
12+
13+
14+
@app.get("/fastapi-info")
15+
def fastapi_info():
16+
return get_fastapi_info()
17+
18+
19+
@app.get("/items/{item_id}")
20+
def read_item(item_id: int, q: Optional[str] = None):
21+
return get_fastapi_info()
22+
23+
24+
@app.exception_handler(StarletteHTTPException)
25+
async def custom_http_exception_handler(request, exc):
26+
return JSONResponse(
27+
status_code=status.HTTP_404_NOT_FOUND,
28+
content=get_fastapi_info(),
29+
)

‎python/sqlcommenter-python/tests/fastapi/tests.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#!/usr/bin/python
2+
#
3+
# Copyright 2019 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import json
18+
19+
import fastapi
20+
import pytest
21+
from starlette.testclient import TestClient
22+
23+
from google.cloud.sqlcommenter.fastapi import get_fastapi_info
24+
25+
from .app import app
26+
27+
28+
@pytest.fixture
29+
def client():
30+
client = TestClient(app)
31+
yield client
32+
33+
34+
def test_get_fastapi_info_in_request_context(client):
35+
expected = {
36+
'app_name': 'SQLCommenter',
37+
'controller': 'fastapi_info',
38+
'framework': 'fastapi:%s' % fastapi.__version__,
39+
'route': '/fastapi-info',
40+
}
41+
resp = client.get('/fastapi-info')
42+
assert json.loads(resp.content.decode('utf-8')) == expected
43+
44+
45+
def test_get_fastapi_info_in_404_error_context(client):
46+
expected = {
47+
'app_name': 'SQLCommenter',
48+
'framework': 'fastapi:%s' % fastapi.__version__,
49+
}
50+
resp = client.get('/doesnt-exist')
51+
assert json.loads(resp.content.decode('utf-8')) == expected
52+
53+
54+
def test_get_fastapi_info_outside_request_context(client):
55+
assert get_fastapi_info() == {}

‎python/sqlcommenter-python/tests/sqlalchemy/tests.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,3 +126,38 @@ def test_route_disabled(self, get_info):
126126
"SELECT 1; /*controller='c',framework='flask'*/",
127127
with_route=False,
128128
)
129+
130+
131+
class FastAPITests(SQLAlchemyTestCase):
132+
fastapi_info = {
133+
'framework': 'fastapi',
134+
'controller': 'c',
135+
'route': '/',
136+
}
137+
138+
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
139+
def test_all_data(self, get_info):
140+
self.assertSQL(
141+
"SELECT 1; /*controller='c',framework='fastapi',route='/'*/",
142+
)
143+
144+
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
145+
def test_framework_disabled(self, get_info):
146+
self.assertSQL(
147+
"SELECT 1; /*controller='c',route='/'*/",
148+
with_framework=False,
149+
)
150+
151+
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
152+
def test_controller_disabled(self, get_info):
153+
self.assertSQL(
154+
"SELECT 1; /*framework='fastapi',route='/'*/",
155+
with_controller=False,
156+
)
157+
158+
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
159+
def test_route_disabled(self, get_info):
160+
self.assertSQL(
161+
"SELECT 1; /*controller='c',framework='fastapi'*/",
162+
with_route=False,
163+
)

0 commit comments

Comments
 (0)