Add FastAPI implementation to OTEL

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.

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

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)}

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)

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

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,
+        )

python/sqlcommenter-python/tox.ini

@@ -24,6 +24,8 @@ deps =
     psycopg2: psycopg2-binary
     sqlalchemy: sqlalchemy
     six
+    asgiref
+    fastapi
 commands =
     python runtests.py