From 2d369897d252ce2a1fda1a7696c607ff53c38924 Mon Sep 17 00:00:00 2001 From: Thomas Wunderlich Date: Thu, 10 Oct 2024 00:41:30 -0400 Subject: [PATCH 1/2] Add FastAPI implementation to OTEL This copies over the existing fastapi implementation from https://github.com/google/sqlcommenter/tree/master --- python/sqlcommenter-python/README.md | 13 +++ .../opentelemetry/sqlcommenter/fastapi.py | 99 +++++++++++++++++++ .../sqlcommenter/sqlalchemy/executor.py | 9 +- .../tests/fastapi/__init__.py | 0 .../sqlcommenter-python/tests/fastapi/app.py | 44 +++++++++ .../tests/fastapi/tests.py | 64 ++++++++++++ .../tests/sqlalchemy/tests.py | 34 +++++++ 7 files changed, 262 insertions(+), 1 deletion(-) create mode 100644 python/sqlcommenter-python/opentelemetry/sqlcommenter/fastapi.py create mode 100644 python/sqlcommenter-python/tests/fastapi/__init__.py create mode 100644 python/sqlcommenter-python/tests/fastapi/app.py create mode 100644 python/sqlcommenter-python/tests/fastapi/tests.py diff --git a/python/sqlcommenter-python/README.md b/python/sqlcommenter-python/README.md index 7a5bcc1..166ef9c 100644 --- a/python/sqlcommenter-python/README.md +++ b/python/sqlcommenter-python/README.md @@ -72,6 +72,19 @@ traceparent='00-5bd66ef5095369c7b0d1f8f4bd33716a-c532cb4098ac3dd2-01', tracestate='congo%%3Dt61rcWkgMzE%%2Crojo%%3D00f067aa0ba902b7'*/ ``` +#### FastAPI + +If you are using SQLAlchemy with FastAPI, add the middleware to get: framework, app_name, controller and route. + +```python +from fastapi import FastAPI +from opentelemetry.sqlcommenter.fastapi import SQLCommenterMiddleware + +app = FastAPI() + +app.add_middleware(SQLCommenterMiddleware) + + ### Psycopg2 Use the provided cursor factory to generate database cursors. All queries executed with such cursors will have the SQL comment prepended to them. diff --git a/python/sqlcommenter-python/opentelemetry/sqlcommenter/fastapi.py b/python/sqlcommenter-python/opentelemetry/sqlcommenter/fastapi.py new file mode 100644 index 0000000..849bad0 --- /dev/null +++ b/python/sqlcommenter-python/opentelemetry/sqlcommenter/fastapi.py @@ -0,0 +1,99 @@ +#!/usr/bin/python +# +# Copyright The OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import + +try: + from typing import Optional + from asgiref.compatibility import guarantee_single_callable + from contextvars import ContextVar + import fastapi + from fastapi import FastAPI + from starlette.routing import Match, Route +except ImportError: + fastapi = None + +context = ContextVar("context", default={}) + + +def get_fastapi_info(): + """ + Get available info from the current FastAPI request, if we're in a + FastAPI request-response cycle. Else, return an empty dict. + """ + info = {} + if fastapi and context: + info = context.get() + return info + + +class SQLCommenterMiddleware: + """The ASGI application middleware. + This class is an ASGI middleware that augment SQL statements before execution, + with comments containing information about the code that caused its execution. + + Args: + app: The ASGI application callable to forward requests to. + """ + + def __init__(self, app): + self.app = guarantee_single_callable(app) + + async def __call__(self, scope, receive, send): + """The ASGI application + Args: + scope: An ASGI environment. + receive: An awaitable callable yielding dictionaries + send: An awaitable callable taking a single dictionary as argument. + """ + if scope["type"] not in ("http", "websocket"): + return await self.app(scope, receive, send) + + if not isinstance(scope["app"], FastAPI): + return await self.app(scope, receive, send) + + fastapi_app = scope["app"] + info = _get_fastapi_info(fastapi_app, scope) + token = context.set(info) + + try: + await self.app(scope, receive, send) + finally: + context.reset(token) + + +def _get_fastapi_info(fastapi_app: FastAPI, scope) -> dict: + info = { + "framework": 'fastapi:%s' % fastapi.__version__, + "app_name": fastapi_app.title, + } + + route = _get_fastapi_route(fastapi_app, scope) + if route: + info["controller"] = route.name + info["route"] = route.path + + return info + + +def _get_fastapi_route(fastapi_app: FastAPI, scope) -> Optional[Route]: + for route in fastapi_app.router.routes: + # Determine if any route matches the incoming scope, + # and return the route name if found. + match, child_scope = route.matches(scope) + if match == Match.FULL: + return child_scope.get("route") + return None diff --git a/python/sqlcommenter-python/opentelemetry/sqlcommenter/sqlalchemy/executor.py b/python/sqlcommenter-python/opentelemetry/sqlcommenter/sqlalchemy/executor.py index 7a218a2..9ba3dee 100644 --- a/python/sqlcommenter-python/opentelemetry/sqlcommenter/sqlalchemy/executor.py +++ b/python/sqlcommenter-python/opentelemetry/sqlcommenter/sqlalchemy/executor.py @@ -22,6 +22,7 @@ import sqlalchemy from opentelemetry.sqlcommenter import generate_sql_comment +from opentelemetry.sqlcommenter.fastapi import get_fastapi_info from opentelemetry.sqlcommenter.flask import get_flask_info from opentelemetry.sqlcommenter.opencensus import get_opencensus_values from opentelemetry.sqlcommenter.opentelemetry import get_opentelemetry_values @@ -42,6 +43,12 @@ def BeforeExecuteFactory( 'db_framework': with_db_framework, } + def get_framework_info(): + info = get_flask_info() + if not info: + info = get_fastapi_info() + return info + def before_cursor_execute(conn, cursor, sql, parameters, context, executemany): data = dict( # TODO: Figure out how to retrieve the exact driver version. @@ -54,7 +61,7 @@ def before_cursor_execute(conn, cursor, sql, parameters, context, executemany): # folks using it in a web framework such as flask will # use it in unison with flask but initialize the parts disjointly, # unlike Django which uses ORMs directly as part of the framework. - data.update(get_flask_info()) + data.update(get_framework_info()) # Filter down to just the requested attributes. data = {k: v for k, v in data.items() if attributes.get(k)} diff --git a/python/sqlcommenter-python/tests/fastapi/__init__.py b/python/sqlcommenter-python/tests/fastapi/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/sqlcommenter-python/tests/fastapi/app.py b/python/sqlcommenter-python/tests/fastapi/app.py new file mode 100644 index 0000000..6bb8862 --- /dev/null +++ b/python/sqlcommenter-python/tests/fastapi/app.py @@ -0,0 +1,44 @@ +from typing import Optional + +from fastapi import FastAPI, status +from fastapi.responses import JSONResponse +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") + +app.add_middleware(SQLCommenterMiddleware) + + +@app.get("/fastapi-info") +def fastapi_info(): + return get_fastapi_info() + + +@app.get("/items/{item_id}") +def read_item(item_id: int, q: Optional[str] = None): + return get_fastapi_info() + + +@app.exception_handler(StarletteHTTPException) +async def custom_http_exception_handler(request, exc): + return JSONResponse( + 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 new file mode 100644 index 0000000..2fb8494 --- /dev/null +++ b/python/sqlcommenter-python/tests/fastapi/tests.py @@ -0,0 +1,64 @@ +#!/usr/bin/python +# +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import fastapi +import pytest +from google.cloud.sqlcommenter.fastapi import get_fastapi_info +from starlette.testclient import TestClient + +from .app import app + + +@pytest.fixture +def client(): + client = TestClient(app) + yield client + + +def test_get_fastapi_info_in_request_context(client): + expected = { + 'app_name': 'SQLCommenter', + 'controller': 'fastapi_info', + 'framework': 'fastapi:%s' % fastapi.__version__, + 'route': '/fastapi-info', + } + resp = client.get('/fastapi-info') + assert json.loads(resp.content.decode('utf-8')) == expected + + +def test_get_fastapi_info_in_404_error_context(client): + expected = { + 'app_name': 'SQLCommenter', + 'framework': 'fastapi:%s' % fastapi.__version__, + } + resp = client.get('/doesnt-exist') + assert json.loads(resp.content.decode('utf-8')) == expected + + +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 diff --git a/python/sqlcommenter-python/tests/sqlalchemy/tests.py b/python/sqlcommenter-python/tests/sqlalchemy/tests.py index 3e0feb4..aeeea3a 100644 --- a/python/sqlcommenter-python/tests/sqlalchemy/tests.py +++ b/python/sqlcommenter-python/tests/sqlalchemy/tests.py @@ -126,3 +126,37 @@ def test_route_disabled(self, get_info): "SELECT 1; /*controller='c',framework='flask'*/", with_route=False, ) + +class FastAPITests(SQLAlchemyTestCase): + fastapi_info = { + 'framework': 'fastapi', + 'controller': 'c', + 'route': '/', + } + + @mock.patch('opentelemetry.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info) + def test_all_data(self, get_info): + self.assertSQL( + "SELECT 1 /*controller='c',framework='fastapi',route='/'*/;", + ) + + @mock.patch('opentelemetry.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info) + def test_framework_disabled(self, get_info): + self.assertSQL( + "SELECT 1 /*controller='c',route='/'*/;", + with_framework=False, + ) + + @mock.patch('opentelemetry.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info) + def test_controller_disabled(self, get_info): + self.assertSQL( + "SELECT 1 /*framework='fastapi',route='/'*/;", + with_controller=False, + ) + + @mock.patch('opentelemetry.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info) + def test_route_disabled(self, get_info): + self.assertSQL( + "SELECT 1 /*controller='c',framework='fastapi'*/;", + with_route=False, + ) \ No newline at end of file From e144e28f20c162ce0f18797d8fc8337b601acfaa Mon Sep 17 00:00:00 2001 From: Thomas Wunderlich Date: Thu, 10 Oct 2024 10:26:33 -0400 Subject: [PATCH 2/2] Add tox config --- python/sqlcommenter-python/tox.ini | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sqlcommenter-python/tox.ini b/python/sqlcommenter-python/tox.ini index 4143314..aa16b0b 100644 --- a/python/sqlcommenter-python/tox.ini +++ b/python/sqlcommenter-python/tox.ini @@ -24,6 +24,8 @@ deps = psycopg2: psycopg2-binary sqlalchemy: sqlalchemy six + asgiref + fastapi commands = python runtests.py