Add DB-API instrumentor support for MySQL driver sqlcommenting (#2897)

tammy-baylis-swi · web-flow · commit a940fc4a65b6 · 2024-10-30T10:46:11.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2922](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2922))
 - `opentelemetry-instrumentation-celery` Don't detach context without a None token
   ([#2927](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2927))
+- `opentelemetry-instrumentation-dbapi` sqlcommenter key values created from PostgreSQL, MySQL systems
+  ([#2897](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2897))
 
 ### Breaking changes
 
diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py b/instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py
@@ -53,6 +53,11 @@
 )
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import SpanKind, TracerProvider, get_tracer
+from opentelemetry.util._importlib_metadata import version as util_version
+
+_DB_DRIVER_ALIASES = {
+    "MySQLdb": "mysqlclient",
+}
 
 _logger = logging.getLogger(__name__)
 
@@ -275,6 +280,70 @@ def __init__(
         self.name = ""
         self.database = ""
         self.connect_module = connect_module
+        self.commenter_data = self.calculate_commenter_data()
+
+    def _get_db_version(
+        self,
+        db_driver,
+    ):
+        if db_driver in _DB_DRIVER_ALIASES:
+            return util_version(_DB_DRIVER_ALIASES[db_driver])
+        db_version = ""
+        try:
+            db_version = self.connect_module.__version__
+        except AttributeError:
+            db_version = "unknown"
+        return db_version
+
+    def calculate_commenter_data(
+        self,
+    ):
+        commenter_data = {}
+        if not self.enable_commenter:
+            return commenter_data
+
+        db_driver = getattr(self.connect_module, "__name__", "unknown")
+        db_version = self._get_db_version(db_driver)
+
+        commenter_data = {
+            "db_driver": f"{db_driver}:{db_version.split(' ')[0]}",
+            # PEP 249-compliant drivers should have the following attributes.
+            # We can assume apilevel "1.0" if not given.
+            # We use "unknown" for others to prevent uncaught AttributeError.
+            # https://peps.python.org/pep-0249/#globals
+            "dbapi_threadsafety": getattr(
+                self.connect_module, "threadsafety", "unknown"
+            ),
+            "dbapi_level": getattr(self.connect_module, "apilevel", "1.0"),
+            "driver_paramstyle": getattr(
+                self.connect_module, "paramstyle", "unknown"
+            ),
+        }
+
+        if self.database_system == "postgresql":
+            if hasattr(self.connect_module, "__libpq_version__"):
+                libpq_version = self.connect_module.__libpq_version__
+            else:
+                libpq_version = self.connect_module.pq.__build_version__
+            commenter_data.update(
+                {
+                    "libpq_version": libpq_version,
+                }
+            )
+        elif self.database_system == "mysql":
+            mysqlc_version = ""
+            if db_driver == "MySQLdb":
+                mysqlc_version = self.connect_module._mysql.get_client_info()
+            elif db_driver == "pymysql":
+                mysqlc_version = self.connect_module.get_client_info()
+
+            commenter_data.update(
+                {
+                    "mysql_client_version": mysqlc_version,
+                }
+            )
+
+        return commenter_data
 
     def wrapped_connection(
         self,
@@ -427,21 +496,23 @@ def traced_execution(
             if args and self._commenter_enabled:
                 try:
                     args_list = list(args)
-                    if hasattr(self._connect_module, "__libpq_version__"):
-                        libpq_version = self._connect_module.__libpq_version__
-                    else:
-                        libpq_version = (
-                            self._connect_module.pq.__build_version__
-                        )
 
-                    commenter_data = {
-                        # Psycopg2/framework information
-                        "db_driver": f"psycopg2:{self._connect_module.__version__.split(' ')[0]}",
-                        "dbapi_threadsafety": self._connect_module.threadsafety,
-                        "dbapi_level": self._connect_module.apilevel,
-                        "libpq_version": libpq_version,
-                        "driver_paramstyle": self._connect_module.paramstyle,
-                    }
+                    # lazy capture of mysql-connector client version using cursor
+                    if (
+                        self._db_api_integration.database_system == "mysql"
+                        and self._db_api_integration.connect_module.__name__
+                        == "mysql.connector"
+                        and not self._db_api_integration.commenter_data[
+                            "mysql_client_version"
+                        ]
+                    ):
+                        self._db_api_integration.commenter_data[
+                            "mysql_client_version"
+                        ] = cursor._cnx._cmysql.get_client_info()
+
+                    commenter_data = dict(
+                        self._db_api_integration.commenter_data
+                    )
                     if self._commenter_options.get(
                         "opentelemetry_values", True
                     ):
diff --git a/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py b/instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py
@@ -24,6 +24,7 @@
 from opentelemetry.test.test_base import TestBase
 
 
+# pylint: disable=too-many-public-methods
 class TestDBApiIntegration(TestBase):
     def setUp(self):
         super().setUp()
@@ -252,6 +253,7 @@ def test_executemany(self):
 
     def test_executemany_comment(self):
         connect_module = mock.MagicMock()
+        connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
         connect_module.__libpq_version__ = 123
         connect_module.apilevel = 123
@@ -260,7 +262,7 @@ def test_executemany_comment(self):
 
         db_integration = dbapi.DatabaseApiIntegration(
             "testname",
-            "testcomponent",
+            "postgresql",
             enable_commenter=True,
             commenter_options={"db_driver": False, "dbapi_level": False},
             connect_module=connect_module,
@@ -275,8 +277,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_comment_non_pep_249_compliant(self):
+        class MockConnectModule:
+            def __getattr__(self, name):
+                if name == "__name__":
+                    return "test"
+                if name == "__version__":
+                    return mock.MagicMock()
+                if name == "__libpq_version__":
+                    return 123
+                raise AttributeError("attribute missing")
+
+        connect_module = MockConnectModule()
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "postgresql",
+            enable_commenter=True,
+            connect_module=connect_module,
+            commenter_options={"db_driver": False},
+        )
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*dbapi_level='1.0',dbapi_threadsafety='unknown',driver_paramstyle='unknown',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
     def test_compatible_build_version_psycopg_psycopg2_libpq(self):
         connect_module = mock.MagicMock()
+        connect_module.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
         connect_module.pq = mock.MagicMock()
         connect_module.pq.__build_version__ = 123
@@ -286,7 +318,7 @@ def test_compatible_build_version_psycopg_psycopg2_libpq(self):
 
         db_integration = dbapi.DatabaseApiIntegration(
             "testname",
-            "testcomponent",
+            "postgresql",
             enable_commenter=True,
             commenter_options={"db_driver": False, "dbapi_level": False},
             connect_module=connect_module,
@@ -301,8 +333,150 @@ def test_compatible_build_version_psycopg_psycopg2_libpq(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_psycopg2_integration_comment(self):
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "psycopg2"
+        connect_module.__version__ = "1.2.3"
+        connect_module.__libpq_version__ = 123
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "postgresql",
+            enable_commenter=True,
+            commenter_options={"db_driver": True, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*db_driver='psycopg2%%3A1.2.3',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_psycopg_integration_comment(self):
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "psycopg"
+        connect_module.__version__ = "1.2.3"
+        connect_module.pq = mock.MagicMock()
+        connect_module.pq.__build_version__ = 123
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "postgresql",
+            enable_commenter=True,
+            commenter_options={"db_driver": True, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*db_driver='psycopg%%3A1.2.3',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_mysqlconnector_integration_comment(self):
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "mysql.connector"
+        connect_module.__version__ = "1.2.3"
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "mysql",
+            enable_commenter=True,
+            commenter_options={"db_driver": True, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*db_driver='mysql.connector%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='1.2.3',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
+    @mock.patch("opentelemetry.instrumentation.dbapi.util_version")
+    def test_executemany_mysqlclient_integration_comment(
+        self,
+        mock_dbapi_util_version,
+    ):
+        mock_dbapi_util_version.return_value = "1.2.3"
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "MySQLdb"
+        connect_module.__version__ = "1.2.3"
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+        connect_module._mysql = mock.MagicMock()
+        connect_module._mysql.get_client_info = mock.MagicMock(
+            return_value="123"
+        )
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "mysql",
+            enable_commenter=True,
+            commenter_options={"db_driver": True, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*db_driver='MySQLdb%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_version='123',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
+    def test_executemany_pymysql_integration_comment(self):
+        connect_module = mock.MagicMock()
+        connect_module.__name__ = "pymysql"
+        connect_module.__version__ = "1.2.3"
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+        connect_module.get_client_info = mock.MagicMock(return_value="123")
+
+        db_integration = dbapi.DatabaseApiIntegration(
+            "testname",
+            "mysql",
+            enable_commenter=True,
+            commenter_options={"db_driver": True, "dbapi_level": False},
+            connect_module=connect_module,
+        )
+
+        mock_connection = db_integration.wrapped_connection(
+            mock_connect, {}, {}
+        )
+        cursor = mock_connection.cursor()
+        cursor.executemany("Select 1;")
+        self.assertRegex(
+            cursor.query,
+            r"Select 1 /\*db_driver='pymysql%%3A1.2.3',dbapi_threadsafety=123,driver_paramstyle='test',mysql_client_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.__name__ = "test"
         connect_module.__version__ = mock.MagicMock()
         connect_module.__libpq_version__ = 123
         connect_module.apilevel = 123
@@ -311,7 +485,7 @@ def test_executemany_flask_integration_comment(self):
 
         db_integration = dbapi.DatabaseApiIntegration(
             "testname",
-            "testcomponent",
+            "postgresql",
             enable_commenter=True,
             commenter_options={"db_driver": False, "dbapi_level": False},
             connect_module=connect_module,
@@ -332,6 +506,11 @@ def test_executemany_flask_integration_comment(self):
             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}'\*/;",
         )
 
+        clear_context = context.set_value(
+            "SQLCOMMENTER_ORM_TAGS_AND_VALUES", {}, current_context
+        )
+        context.attach(clear_context)
+
     def test_callproc(self):
         db_integration = dbapi.DatabaseApiIntegration(
             "testname", "testcomponent"
@@ -415,6 +594,12 @@ class MockCursor:
     def __init__(self) -> None:
         self.query = ""
         self.params = None
+        # Mock mysql.connector modules and method
+        self._cnx = mock.MagicMock()
+        self._cnx._cmysql = mock.MagicMock()
+        self._cnx._cmysql.get_client_info = mock.MagicMock(
+            return_value="1.2.3"
+        )
 
     # pylint: disable=unused-argument, no-self-use
     def execute(self, query, params=None, throw_exception=False):
