ext/pymysql: Implement instrument_connection, uninstrument_connection

Author: mauriciovasquezbernal

ext/opentelemetry-ext-docker-tests/tests/pymysql/test_pymysql_functional.py

@@ -48,6 +48,7 @@ def setUpClass(cls):
     def tearDownClass(cls):
         if cls._connection:
             cls._connection.close()
+        PyMySQLInstrumentor().uninstrument()
 
     def validate_spans(self):
         spans = self.memory_exporter.get_finished_spans()

ext/opentelemetry-ext-pymysql/src/opentelemetry/ext/pymysql/__init__.py

@@ -48,12 +48,22 @@
 import pymysql
 
 from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.ext.dbapi import unwrap_connect, wrap_connect
+from opentelemetry.ext import dbapi
 from opentelemetry.ext.pymysql.version import __version__
 from opentelemetry.trace import TracerProvider, get_tracer
 
 
 class PyMySQLInstrumentor(BaseInstrumentor):
+    _CONNECTION_ATTRIBUTES = {
+        "database": "db",
+        "port": "port",
+        "host": "host",
+        "user": "user",
+    }
+
+    _DATABASE_COMPONENT = "mysql"
+    _DATABASE_TYPE = "sql"
+
     def _instrument(self, **kwargs):
         """Integrate with the PyMySQL library.
         https://github.com/PyMySQL/PyMySQL/
@@ -62,15 +72,40 @@ def _instrument(self, **kwargs):
 
         tracer = get_tracer(__name__, __version__, tracer_provider)
 
-        connection_attributes = {
-            "database": "db",
-            "port": "port",
-            "host": "host",
-            "user": "user",
-        }
-        wrap_connect(
-            tracer, pymysql, "connect", "mysql", "sql", connection_attributes
+        dbapi.wrap_connect(
+            tracer,
+            pymysql,
+            "connect",
+            self._DATABASE_COMPONENT,
+            self._DATABASE_TYPE,
+            self._CONNECTION_ATTRIBUTES,
         )
 
     def _uninstrument(self, **kwargs):
-        unwrap_connect(pymysql, "connect")
+        """"Disable PyMySQL instrumentation"""
+        dbapi.unwrap_connect(pymysql, "connect")
+
+    # pylint:disable=no-self-use
+    def instrument_connection(self, connection):
+        """Enable instrumentation in a PyMySQL connection.
+
+            Args:
+                connection: The connection to instrument.
+        """
+        tracer = get_tracer(__name__, __version__)
+
+        return dbapi.instrument_connection(
+            tracer,
+            connection,
+            self._DATABASE_COMPONENT,
+            self._DATABASE_TYPE,
+            self._CONNECTION_ATTRIBUTES,
+        )
+
+    def uninstrument_connection(self, connection):
+        """Disable instrumentation in a PyMySQL connection.
+
+            Args:
+                connection: The connection to uninstrument.
+        """
+        return dbapi.uninstrument_connection(connection)

ext/opentelemetry-ext-pymysql/tests/test_pymysql_integration.py

@@ -75,3 +75,40 @@ def test_custom_tracer_provider(self, mock_connect):
         span = spans_list[0]
 
         self.assertIs(span.resource, resource)
+
+    @mock.patch("pymysql.connect")
+    # pylint: disable=unused-argument
+    def test_instrument_connection(self, mock_connect):
+        cnx = pymysql.connect(database="test")
+        query = "SELECT * FROM test"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 0)
+
+        cnx = PyMySQLInstrumentor().instrument_connection(cnx)
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+
+    @mock.patch("pymysql.connect")
+    # pylint: disable=unused-argument
+    def test_uninstrument_connection(self, mock_connect):
+        PyMySQLInstrumentor().instrument()
+        cnx = pymysql.connect(database="test")
+        query = "SELECT * FROM test"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+
+        cnx = PyMySQLInstrumentor().uninstrument_connection(cnx)
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)