Flask sqlalchemy psycopg2 integration (#1224)

Author: Thiyagu55

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Flask sqlalchemy psycopg2 integration
+  ([#1224](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1224))
+
 ### Fixed
 
 - `opentelemetry-instrumentation-boto3sqs` Make propagation compatible with other SQS instrumentations, add 'messaging.url' span attribute, and fix missing package dependencies.

instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py

@@ -45,8 +45,8 @@
 
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.dbapi.version import __version__
+from opentelemetry.instrumentation.sqlcommenter_utils import _add_sql_comment
 from opentelemetry.instrumentation.utils import (
-    _add_sql_comment,
     _get_opentelemetry_values,
     unwrap,
 )

instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py

@@ -16,6 +16,7 @@
 import logging
 from unittest import mock
 
+from opentelemetry import context
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation import dbapi
 from opentelemetry.sdk import resources
@@ -254,6 +255,38 @@ def test_executemany_comment(self):
             r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
 
+    def test_executemany_flask_integration_comment(self):
+
+        connect_module = mock.MagicMock()
+        connect_module.__version__ = mock.MagicMock()
+        connect_module.__libpq_version__ = 123
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "testcomponent",
+            enable_commenter=True,
+            commenter_options={"db_driver": False, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+        current_context = context.get_current()
+        sqlcommenter_context = context.set_value(
+            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
+        )
+        context.attach(sqlcommenter_context)
+
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',flask=1,libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
     def test_callproc(self):
         db_integration = dbapi.DatabaseApiIntegration(
             "testname", "testcomponent"

instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware/sqlcommenter_middleware.py

@@ -22,10 +22,8 @@
 from django.db import connections
 from django.db.backends.utils import CursorDebugWrapper
 
-from opentelemetry.instrumentation.utils import (
-    _add_sql_comment,
-    _get_opentelemetry_values,
-)
+from opentelemetry.instrumentation.sqlcommenter_utils import _add_sql_comment
+from opentelemetry.instrumentation.utils import _get_opentelemetry_values
 from opentelemetry.trace.propagation.tracecontext import (
     TraceContextTextMapPropagator,
 )

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -24,6 +24,53 @@
 * The ``http.route`` Span attribute is set so that one can see which URL rule
   matched a request.
 
+SQLCOMMENTER
+*****************************************
+You can optionally configure Flask instrumentation to enable sqlcommenter which enriches
+the query with contextual information.
+
+Usage
+-----
+
+.. code:: python
+
+    from opentelemetry.instrumentation.flask import FlaskInstrumentor
+
+    FlaskInstrumentor().instrument(enable_commenter=True, commenter_options={})
+
+
+For example,
+::
+
+   FlaskInstrumentor when used with SQLAlchemyInstrumentor or Psycopg2Instrumentor, invoking cursor.execute("select * from auth_users")
+   will lead to sql query "select * from auth_users" but when SQLCommenter is enabled
+   the query will get appended with some configurable tags like "select * from auth_users /*metrics=value*/;"
+
+    Inorder for the commenter to append flask related tags to sql queries, the commenter needs to enabled on
+    the respective  SQLAlchemyInstrumentor or Psycopg2Instrumentor framework too.
+
+SQLCommenter Configurations
+***************************
+We can configure the tags to be appended to the sqlquery log by adding configuration inside commenter_options(default:{}) keyword
+
+framework = True(Default) or False
+
+For example,
+::
+Enabling this flag will add flask and it's version which is /*flask%%3A2.9.3*/
+
+route = True(Default) or False
+
+For example,
+::
+Enabling this flag will add route uri /*route='/home'*/
+
+controller = True(Default) or False
+
+For example,
+::
+Enabling this flag will add controller name /*controller='home_view'*/
+
 Usage
 -----
 
@@ -255,6 +302,8 @@ def _wrapped_before_request(
     request_hook=None,
     tracer=None,
     excluded_urls=None,
+    enable_commenter=True,
+    commenter_options=None,
 ):
     def _before_request():
         if excluded_urls and excluded_urls.url_disabled(flask.request.url):
@@ -300,6 +349,30 @@ def _before_request():
         flask_request_environ[_ENVIRON_SPAN_KEY] = span
         flask_request_environ[_ENVIRON_TOKEN] = token
 
+        if enable_commenter:
+            current_context = context.get_current()
+            flask_info = {}
+
+            # https://flask.palletsprojects.com/en/1.1.x/api/#flask.has_request_context
+            if flask and flask.request:
+                if commenter_options.get("framework", True):
+                    flask_info["framework"] = f"flask:{flask.__version__}"
+                if (
+                    commenter_options.get("controller", True)
+                    and flask.request.endpoint
+                ):
+                    flask_info["controller"] = flask.request.endpoint
+                if (
+                    commenter_options.get("route", True)
+                    and flask.request.url_rule
+                    and flask.request.url_rule.rule
+                ):
+                    flask_info["route"] = flask.request.url_rule.rule
+            sqlcommenter_context = context.set_value(
+                "SQLCOMMENTER_ORM_TAGS_AND_VALUES", flask_info, current_context
+            )
+            context.attach(sqlcommenter_context)
+
     return _before_request
 
 
@@ -336,6 +409,8 @@ class _InstrumentedFlask(flask.Flask):
     _tracer_provider = None
     _request_hook = None
     _response_hook = None
+    _enable_commenter = True
+    _commenter_options = None
     _meter_provider = None
 
     def __init__(self, *args, **kwargs):
@@ -374,6 +449,8 @@ def __init__(self, *args, **kwargs):
             _InstrumentedFlask._request_hook,
             tracer,
             excluded_urls=_InstrumentedFlask._excluded_urls,
+            enable_commenter=_InstrumentedFlask._enable_commenter,
+            commenter_options=_InstrumentedFlask._commenter_options,
         )
         self._before_request = _before_request
         self.before_request(_before_request)
@@ -410,6 +487,11 @@ def _instrument(self, **kwargs):
             if excluded_urls is None
             else parse_excluded_urls(excluded_urls)
         )
+        enable_commenter = kwargs.get("enable_commenter", True)
+        _InstrumentedFlask._enable_commenter = enable_commenter
+
+        commenter_options = kwargs.get("commenter_options", {})
+        _InstrumentedFlask._commenter_options = commenter_options
         meter_provider = kwargs.get("meter_provider")
         _InstrumentedFlask._meter_provider = meter_provider
         flask.Flask = _InstrumentedFlask
@@ -424,6 +506,8 @@ def instrument_app(
         response_hook=None,
         tracer_provider=None,
         excluded_urls=None,
+        enable_commenter=True,
+        commenter_options=None,
         meter_provider=None,
     ):
         if not hasattr(app, "_is_instrumented_by_opentelemetry"):
@@ -462,6 +546,10 @@ def instrument_app(
                 request_hook,
                 tracer,
                 excluded_urls=excluded_urls,
+                enable_commenter=enable_commenter,
+                commenter_options=commenter_options
+                if commenter_options
+                else {},
             )
             app._before_request = _before_request
             app.before_request(_before_request)

instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py

@@ -16,6 +16,8 @@
 from werkzeug.test import Client
 from werkzeug.wrappers import Response
 
+from opentelemetry import context
+
 
 class InstrumentationTest:
     @staticmethod
@@ -24,6 +26,14 @@ def _hello_endpoint(helloid):
             raise ValueError(":-(")
         return "Hello: " + str(helloid)
 
+    @staticmethod
+    def _sqlcommenter_endpoint():
+        current_context = context.get_current()
+        sqlcommenter_flask_values = current_context.get(
+            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}
+        )
+        return sqlcommenter_flask_values
+
     @staticmethod
     def _custom_response_headers():
         resp = flask.Response("test response")
@@ -43,6 +53,7 @@ def excluded2_endpoint():
 
         # pylint: disable=no-member
         self.app.route("/hello/<int:helloid>")(self._hello_endpoint)
+        self.app.route("/sqlcommenter")(self._sqlcommenter_endpoint)
         self.app.route("/excluded/<int:helloid>")(self._hello_endpoint)
         self.app.route("/excluded")(excluded_endpoint)
         self.app.route("/excluded2")(excluded2_endpoint)

instrumentation/opentelemetry-instrumentation-flask/tests/test_sqlcommenter.py

@@ -0,0 +1,78 @@
+# 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.
+
+import flask
+from werkzeug.test import Client
+from werkzeug.wrappers import Response
+
+from opentelemetry.instrumentation.flask import FlaskInstrumentor
+from opentelemetry.test.wsgitestutil import WsgiTestBase
+
+# pylint: disable=import-error
+from .base_test import InstrumentationTest
+
+
+class TestSQLCommenter(InstrumentationTest, WsgiTestBase):
+    def setUp(self):
+        super().setUp()
+        FlaskInstrumentor().instrument()
+        self.app = flask.Flask(__name__)
+        self._common_initialization()
+
+    def tearDown(self):
+        super().tearDown()
+        with self.disable_logging():
+            FlaskInstrumentor().uninstrument()
+
+    def test_sqlcommenter_enabled_default(self):
+
+        self.app = flask.Flask(__name__)
+        self.app.route("/sqlcommenter")(self._sqlcommenter_endpoint)
+        client = Client(self.app, Response)
+
+        resp = client.get("/sqlcommenter")
+        self.assertEqual(200, resp.status_code)
+        self.assertRegex(
+            list(resp.response)[0].strip(),
+            b'{"controller":"_sqlcommenter_endpoint","framework":"flask:(.*)","route":"/sqlcommenter"}',
+        )
+
+    def test_sqlcommenter_enabled_with_configurations(self):
+        FlaskInstrumentor().uninstrument()
+        FlaskInstrumentor().instrument(
+            enable_commenter=True, commenter_options={"route": False}
+        )
+
+        self.app = flask.Flask(__name__)
+        self.app.route("/sqlcommenter")(self._sqlcommenter_endpoint)
+        client = Client(self.app, Response)
+
+        resp = client.get("/sqlcommenter")
+        self.assertEqual(200, resp.status_code)
+        self.assertRegex(
+            list(resp.response)[0].strip(),
+            b'{"controller":"_sqlcommenter_endpoint","framework":"flask:(.*)"}',
+        )
+
+    def test_sqlcommenter_disabled(self):
+        FlaskInstrumentor().uninstrument()
+        FlaskInstrumentor().instrument(enable_commenter=False)
+
+        self.app = flask.Flask(__name__)
+        self.app.route("/sqlcommenter")(self._sqlcommenter_endpoint)
+        client = Client(self.app, Response)
+
+        resp = client.get("/sqlcommenter")
+        self.assertEqual(200, resp.status_code)
+        self.assertEqual(list(resp.response)[0].strip(), b"{}")

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

@@ -20,10 +20,8 @@
     _instrumenting_module_name,
 )
 from opentelemetry.instrumentation.sqlalchemy.version import __version__
-from opentelemetry.instrumentation.utils import (
-    _add_sql_comment,
-    _get_opentelemetry_values,
-)
+from opentelemetry.instrumentation.sqlcommenter_utils import _add_sql_comment
+from opentelemetry.instrumentation.utils import _get_opentelemetry_values
 from opentelemetry.semconv.trace import NetTransportValues, SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
 

instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py

@@ -16,6 +16,7 @@
 import pytest
 from sqlalchemy import create_engine
 
+from opentelemetry import context
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 from opentelemetry.test.test_base import TestBase
 
@@ -54,3 +55,25 @@ def test_sqlcommenter_enabled(self):
             self.caplog.records[-2].getMessage(),
             r"SELECT  1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
+
+    def test_sqlcommenter_flask_integration(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+            commenter_options={"db_framework": False},
+        )
+        cnx = engine.connect()
+
+        current_context = context.get_current()
+        sqlcommenter_context = context.set_value(
+            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
+        )
+        context.attach(sqlcommenter_context)
+
+        cnx.execute("SELECT  1;").fetchall()
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1 /\*db_driver='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )