ext/pymysql: Add Instrumentor (#611)

Co-authored-by: Diego Hurtado <[email protected]>

Author: mauriciovasquezbernal

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

@@ -125,6 +125,14 @@ def wrap_connect_(
         logger.warning("Failed to integrate with DB API. %s", str(ex))
 
 
+def unwrap_connect(
+    connect_module: typing.Callable[..., any], connect_method_name: str,
+):
+    conn = getattr(connect_module, connect_method_name, None)
+    if isinstance(conn, wrapt.ObjectProxy):
+        setattr(connect_module, connect_method_name, conn.__wrapped__)
+
+
 class DatabaseApiIntegration:
     def __init__(
         self,
@@ -184,7 +192,7 @@ def get_connection_attributes(self, connection):
             self.name += "." + self.database
         user = self.connection_props.get("user")
         if user is not None:
-            self.span_attributes["db.user"] = user
+            self.span_attributes["db.user"] = str(user)
         host = self.connection_props.get("host")
         if host is not None:
             self.span_attributes["net.peer.name"] = host

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

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest import mock
+
 from opentelemetry import trace as trace_api
-from opentelemetry.ext.dbapi import DatabaseApiIntegration
+from opentelemetry.ext import dbapi
 from opentelemetry.test.test_base import TestBase
 
 
@@ -35,7 +37,7 @@ def test_span_succeeded(self):
             "host": "server_host",
             "user": "user",
         }
-        db_integration = DatabaseApiIntegration(
+        db_integration = dbapi.DatabaseApiIntegration(
             self.tracer, "testcomponent", "testtype", connection_attributes
         )
         mock_connection = db_integration.wrapped_connection(
@@ -66,7 +68,9 @@ def test_span_succeeded(self):
         )
 
     def test_span_failed(self):
-        db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
+        db_integration = dbapi.DatabaseApiIntegration(
+            self.tracer, "testcomponent"
+        )
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
         )
@@ -85,7 +89,9 @@ def test_span_failed(self):
         self.assertEqual(span.status.description, "Test Exception")
 
     def test_executemany(self):
-        db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
+        db_integration = dbapi.DatabaseApiIntegration(
+            self.tracer, "testcomponent"
+        )
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
         )
@@ -97,7 +103,9 @@ def test_executemany(self):
         self.assertEqual(span.attributes["db.statement"], "Test query")
 
     def test_callproc(self):
-        db_integration = DatabaseApiIntegration(self.tracer, "testcomponent")
+        db_integration = dbapi.DatabaseApiIntegration(
+            self.tracer, "testcomponent"
+        )
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
         )
@@ -110,6 +118,26 @@ def test_callproc(self):
             span.attributes["db.statement"], "Test stored procedure"
         )
 
+    @mock.patch("opentelemetry.ext.dbapi")
+    def test_wrap_connect(self, mock_dbapi):
+        dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
+        connection = mock_dbapi.connect()
+        self.assertEqual(mock_dbapi.connect.call_count, 1)
+        self.assertIsInstance(connection, dbapi.TracedConnectionProxy)
+        self.assertIsInstance(connection.__wrapped__, mock.Mock)
+
+    @mock.patch("opentelemetry.ext.dbapi")
+    def test_unwrap_connect(self, mock_dbapi):
+        dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
+        connection = mock_dbapi.connect()
+        self.assertEqual(mock_dbapi.connect.call_count, 1)
+        self.assertIsInstance(connection, dbapi.TracedConnectionProxy)
+
+        dbapi.unwrap_connect(mock_dbapi, "connect")
+        connection = mock_dbapi.connect()
+        self.assertEqual(mock_dbapi.connect.call_count, 2)
+        self.assertIsInstance(connection, mock.Mock)
+
 
 # pylint: disable=unused-argument
 def mock_connect(*args, **kwargs):

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

@@ -13,18 +13,12 @@
 # limitations under the License.
 
 import os
-import time
-import unittest
 
 import pymysql as pymy
 
 from opentelemetry import trace as trace_api
-from opentelemetry.ext.pymysql import trace_integration
-from opentelemetry.sdk.trace import Tracer, TracerProvider
-from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
-    InMemorySpanExporter,
-)
+from opentelemetry.ext.pymysql import PyMySQLInstrumentor
+from opentelemetry.test.test_base import TestBase
 
 MYSQL_USER = os.getenv("MYSQL_USER ", "testuser")
 MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD ", "testpassword")
@@ -33,17 +27,14 @@
 MYSQL_DB_NAME = os.getenv("MYSQL_DB_NAME ", "opentelemetry-tests")
 
 
-class TestFunctionalPyMysql(unittest.TestCase):
+class TestFunctionalPyMysql(TestBase):
     @classmethod
     def setUpClass(cls):
+        super().setUpClass()
         cls._connection = None
         cls._cursor = None
-        cls._tracer_provider = TracerProvider()
-        cls._tracer = Tracer(cls._tracer_provider, None)
-        cls._span_exporter = InMemorySpanExporter()
-        cls._span_processor = SimpleExportSpanProcessor(cls._span_exporter)
-        cls._tracer_provider.add_span_processor(cls._span_processor)
-        trace_integration(cls._tracer_provider)
+        cls._tracer = cls.tracer_provider.get_tracer(__name__)
+        PyMySQLInstrumentor().instrument()
         cls._connection = pymy.connect(
             user=MYSQL_USER,
             password=MYSQL_PASSWORD,
@@ -58,11 +49,8 @@ def tearDownClass(cls):
         if cls._connection:
             cls._connection.close()
 
-    def setUp(self):
-        self._span_exporter.clear()
-
     def validate_spans(self):
-        spans = self._span_exporter.get_finished_spans()
+        spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 2)
         for span in spans:
             if span.name == "rootSpan":

ext/opentelemetry-ext-pymysql/CHANGELOG.md

@@ -1,3 +1,6 @@
 # Changelog
 
 ## Unreleased
+
+- Implement PyMySQL integration ([#504](https://github.com/open-telemetry/opentelemetry-python/pull/504))
+- Implement instrumentor interface ([#611](https://github.com/open-telemetry/opentelemetry-python/pull/611))

ext/opentelemetry-ext-pymysql/README.rst

@@ -6,10 +6,6 @@ OpenTelemetry PyMySQL integration
 .. |pypi| image:: https://badge.fury.io/py/opentelemetry-ext-pymysql.svg
    :target: https://pypi.org/project/opentelemetry-ext-pymysql/
 
-Integration with PyMySQL that supports the PyMySQL library and is
-specified to trace_integration using 'PyMySQL'.
-
-
 Installation
 ------------
 

ext/opentelemetry-ext-pymysql/setup.cfg

@@ -42,7 +42,16 @@ packages=find_namespace:
 install_requires =
     opentelemetry-api == 0.7.dev0
     opentelemetry-ext-dbapi == 0.7.dev0
+    opentelemetry-auto-instrumentation == 0.7.dev0
     PyMySQL ~=  0.9.3
 
+[options.extras_require]
+test =
+    opentelemetry-test == 0.7.dev0
+
 [options.packages.find]
 where = src
+
+[options.entry_points]
+opentelemetry_instrumentor =
+    pymysql = opentelemetry.ext.pymysql:PyMySQLInstrumentor

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

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 """
-The integration with PyMySQL supports the `PyMySQL`_ library and is specified
-to ``trace_integration`` using ``'PyMySQL'``.
+The integration with PyMySQL supports the `PyMySQL`_ library and can be enabled
+by using ``PyMySQLInstrumentor``.
 
 .. _PyMySQL: https://pypi.org/project/PyMySQL/
 
@@ -25,11 +25,13 @@
 
     import pymysql
     from opentelemetry import trace
-    from opentelemetry.ext.pymysql import trace_integration
+    from opentelemetry.ext.pymysql import PyMySQLInstrumentor
     from opentelemetry.sdk.trace import TracerProvider
 
     trace.set_tracer_provider(TracerProvider())
-    trace_integration()
+
+    PyMySQLInstrumentor().instrument()
+
     cnx = pymysql.connect(database="MySQL_Database")
     cursor = cnx.cursor()
     cursor.execute("INSERT INTO test (testField) VALUES (123)"
@@ -45,24 +47,30 @@
 
 import pymysql
 
-from opentelemetry.ext.dbapi import wrap_connect
+from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.ext.dbapi import unwrap_connect, wrap_connect
 from opentelemetry.ext.pymysql.version import __version__
 from opentelemetry.trace import TracerProvider, get_tracer
 
 
-def trace_integration(tracer_provider: typing.Optional[TracerProvider] = None):
-    """Integrate with the PyMySQL library.
-       https://github.com/PyMySQL/PyMySQL/
-    """
+class PyMySQLInstrumentor(BaseInstrumentor):
+    def _instrument(self, **kwargs):
+        """Integrate with the PyMySQL library.
+        https://github.com/PyMySQL/PyMySQL/
+        """
+        tracer_provider = kwargs.get("tracer_provider")
+
+        tracer = get_tracer(__name__, __version__, tracer_provider)
 
-    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
+        )
 
-    connection_attributes = {
-        "database": "db",
-        "port": "port",
-        "host": "host",
-        "user": "user",
-    }
-    wrap_connect(
-        tracer, pymysql, "connect", "mysql", "sql", connection_attributes
-    )
+    def _uninstrument(self, **kwargs):
+        unwrap_connect(pymysql, "connect")

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

@@ -17,22 +17,61 @@
 import pymysql
 
 import opentelemetry.ext.pymysql
-from opentelemetry.ext.pymysql import trace_integration
+from opentelemetry.ext.pymysql import PyMySQLInstrumentor
+from opentelemetry.sdk import resources
 from opentelemetry.test.test_base import TestBase
 
 
 class TestPyMysqlIntegration(TestBase):
-    def test_trace_integration(self):
-        with mock.patch("pymysql.connect"):
-            trace_integration()
-            cnx = pymysql.connect(database="test")
-            cursor = cnx.cursor()
-            query = "SELECT * FROM test"
-            cursor.execute(query)
+    def tearDown(self):
+        super().tearDown()
+        with self.disable_logging():
+            PyMySQLInstrumentor().uninstrument()
+
+    @mock.patch("pymysql.connect")
+    # pylint: disable=unused-argument
+    def test_instrumentor(self, mock_connect):
+        PyMySQLInstrumentor().instrument()
+
+        cnx = pymysql.connect(database="test")
+        cursor = cnx.cursor()
+        query = "SELECT * FROM test"
+        cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
         span = spans_list[0]
 
         # Check version and name in span's instrumentation info
         self.check_span_instrumentation_info(span, opentelemetry.ext.pymysql)
+
+        # check that no spans are generated after uninstrument
+        PyMySQLInstrumentor().uninstrument()
+
+        cnx = pymysql.connect(database="test")
+        cursor = cnx.cursor()
+        query = "SELECT * FROM test"
+        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_custom_tracer_provider(self, mock_connect):
+        resource = resources.Resource.create({})
+        result = self.create_tracer_provider(resource=resource)
+        tracer_provider, exporter = result
+
+        PyMySQLInstrumentor().instrument(tracer_provider=tracer_provider)
+
+        cnx = pymysql.connect(database="test")
+        cursor = cnx.cursor()
+        query = "SELECT * FROM test"
+        cursor.execute(query)
+
+        spans_list = exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        self.assertIs(span.resource, resource)

tests/util/src/opentelemetry/test/test_base.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import unittest
+from contextlib import contextmanager
 
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.trace import TracerProvider, export
@@ -59,3 +61,13 @@ def create_tracer_provider(**kwargs):
         tracer_provider.add_span_processor(span_processor)
 
         return tracer_provider, memory_exporter
+
+    @staticmethod
+    @contextmanager
+    def disable_logging(highest_level=logging.CRITICAL):
+        logging.disable(highest_level)
+
+        try:
+            yield
+        finally:
+            logging.disable(logging.NOTSET)

tox.ini

@@ -188,8 +188,9 @@ commands_pre =
   psycopg2: pip install {toxinidir}/ext/opentelemetry-ext-dbapi
   psycopg2: pip install {toxinidir}/ext/opentelemetry-ext-psycopg2
 
+  pymysql: pip install {toxinidir}/opentelemetry-auto-instrumentation
   pymysql: pip install {toxinidir}/ext/opentelemetry-ext-dbapi
-  pymysql: pip install {toxinidir}/ext/opentelemetry-ext-pymysql
+  pymysql: pip install {toxinidir}/ext/opentelemetry-ext-pymysql[test]
 
   redis: pip install {toxinidir}/opentelemetry-auto-instrumentation
   redis: pip install {toxinidir}/ext/opentelemetry-ext-redis[test]