ext/mysql: Add instrumentor interface

ext/opentelemetry-ext-dbapi/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Implement instrument_connection and uninstrument_connection ([#624](https://github.com/open-telemetry/opentelemetry-python/pull/624))
+
 ## 0.4a0
 
 Released 2020-02-21

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

@@ -103,7 +103,7 @@ def wrap_connect(
     """
 
     # pylint: disable=unused-argument
-    def wrap_connect_(
+    def _wrap_connect(
         wrapped: typing.Callable[..., any],
         instance: typing.Any,
         args: typing.Tuple[any, any],
@@ -119,7 +119,7 @@ def wrap_connect_(
 
     try:
         wrapt.wrap_function_wrapper(
-            connect_module, connect_method_name, wrap_connect_
+            connect_module, connect_method_name, _wrap_connect
         )
     except Exception as ex:  # pylint: disable=broad-except
         logger.warning("Failed to integrate with DB API. %s", str(ex))
@@ -128,11 +128,65 @@ def wrap_connect_(
 def unwrap_connect(
     connect_module: typing.Callable[..., any], connect_method_name: str,
 ):
+    """Disable integration with DB API library.
+        https://www.python.org/dev/peps/pep-0249/
+
+        Args:
+            connect_module: Module name where the connect method is available.
+            connect_method_name: The connect method name.
+    """
     conn = getattr(connect_module, connect_method_name, None)
     if isinstance(conn, wrapt.ObjectProxy):
         setattr(connect_module, connect_method_name, conn.__wrapped__)
 
 
+def instrument_connection(
+    tracer,
+    connection,
+    database_component: str,
+    database_type: str = "",
+    connection_attributes: typing.Dict = None,
+):
+    """Enable instrumentation in a database connection.
+
+    Args:
+        tracer: The :class:`Tracer` to use.
+        connection: The connection to instrument.
+        database_component: Database driver name or database name "JDBI",
+            "jdbc", "odbc", "postgreSQL".
+        database_type: The Database type. For any SQL database, "sql".
+        connection_attributes: Attribute names for database, port, host and
+            user in a connection object.
+
+    Returns:
+        An instrumented connection.
+    """
+    db_integration = DatabaseApiIntegration(
+        tracer,
+        database_component,
+        database_type,
+        connection_attributes=connection_attributes,
+    )
+    db_integration.get_connection_attributes(connection)
+    return TracedConnectionProxy(connection, db_integration)
+
+
+def uninstrument_connection(connection):
+    """Disable instrumentation in a database connection.
+
+    Args:
+        connection: The connection to uninstrument.
+
+    Returns:
+        An uninstrumented connection.
+    """
+    if isinstance(connection, wrapt.ObjectProxy):
+        return connection.__wrapped__
+
+    logger.warning("Connection is not instrumented")
+    return connection
+
+
 class DatabaseApiIntegration:
     def __init__(
         self,
@@ -167,8 +221,7 @@ def wrapped_connection(
         """
         connection = connect_method(*args, **kwargs)
         self.get_connection_attributes(connection)
-        traced_connection = TracedConnectionProxy(connection, self)
-        return traced_connection
+        return TracedConnectionProxy(connection, self)
 
     def get_connection_attributes(self, connection):
         # Populate span fields using connection

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

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import logging
 from unittest import mock
 
 from opentelemetry import trace as trace_api
@@ -138,6 +140,30 @@ def test_unwrap_connect(self, mock_dbapi):
         self.assertEqual(mock_dbapi.connect.call_count, 2)
         self.assertIsInstance(connection, mock.Mock)
 
+    def test_instrument_connection(self):
+        connection = mock.Mock()
+        # Avoid get_attributes failing because can't concatenate mock
+        connection.database = "-"
+        connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
+        self.assertIsInstance(connection2, dbapi.TracedConnectionProxy)
+        self.assertIs(connection2.__wrapped__, connection)
+
+    def test_uninstrument_connection(self):
+        connection = mock.Mock()
+        # Set connection.database to avoid a failure because mock can't
+        # be concatenated
+        connection.database = "-"
+        connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
+        self.assertIsInstance(connection2, dbapi.TracedConnectionProxy)
+        self.assertIs(connection2.__wrapped__, connection)
+
+        connection3 = dbapi.uninstrument_connection(connection2)
+        self.assertIs(connection3, connection)
+
+        with self.assertLogs(level=logging.WARNING):
+            connection4 = dbapi.uninstrument_connection(connection)
+        self.assertIs(connection4, connection)
+
 
 # pylint: disable=unused-argument
 def mock_connect(*args, **kwargs):

ext/opentelemetry-ext-docker-tests/tests/mysql/test_mysql_functional.py

@@ -18,7 +18,7 @@
 import mysql.connector
 
 from opentelemetry import trace as trace_api
-from opentelemetry.ext.mysql import trace_integration
+from opentelemetry.ext.mysql import MySQLInstrumentor
 from opentelemetry.test.test_base import TestBase
 
 MYSQL_USER = os.getenv("MYSQL_USER ", "testuser")
@@ -35,7 +35,7 @@ def setUpClass(cls):
         cls._connection = None
         cls._cursor = None
         cls._tracer = cls.tracer_provider.get_tracer(__name__)
-        trace_integration(cls.tracer_provider)
+        MySQLInstrumentor().instrument()
         cls._connection = mysql.connector.connect(
             user=MYSQL_USER,
             password=MYSQL_PASSWORD,
@@ -49,6 +49,7 @@ def setUpClass(cls):
     def tearDownClass(cls):
         if cls._connection:
             cls._connection.close()
+        MySQLInstrumentor().uninstrument()
 
     def validate_spans(self):
         spans = self.memory_exporter.get_finished_spans()

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-mysql/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Implement instrumentor interface ([#654](https://github.com/open-telemetry/opentelemetry-python/pull/654))
+
 ## 0.4a0
 
 Released 2020-02-21

ext/opentelemetry-ext-mysql/setup.cfg

@@ -42,6 +42,7 @@ packages=find_namespace:
 install_requires =
     opentelemetry-api == 0.7.dev0
     opentelemetry-ext-dbapi == 0.7.dev0
+    opentelemetry-auto-instrumentation == 0.7.dev0
     mysql-connector-python ~= 8.0
     wrapt >= 1.0.0, < 2.0.0
 
@@ -51,3 +52,7 @@ test =
 
 [options.packages.find]
 where = src
+
+[options.entry_points]
+opentelemetry_instrumentor =
+    mysql = opentelemetry.ext.pymysql:MySQLInstrumentor

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

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 """
-The integration with MySQL supports the `mysql-connector`_ library and is specified
-to ``trace_integration`` using ``'MySQL'``.
+MySQL instrumentation supporting `mysql-connector`_, it can be enabled by
+using ``MySQLInstrumentor``.
 
 .. _mysql-connector: https://pypi.org/project/mysql-connector/
 
@@ -26,12 +26,13 @@
     import mysql.connector
     from opentelemetry import trace
     from opentelemetry.trace import TracerProvider
-    from opentelemetry.ext.mysql import trace_integration
+    from opentelemetry.ext.mysql import MySQLInstrumentor
 
     trace.set_tracer_provider(TracerProvider())
 
-    trace_integration()
-    cnx = mysql.connector.connect(database='MySQL_Database')
+    MySQLInstrumentor().instrument()
+
+    cnx = mysql.connector.connect(database="MySQL_Database")
     cursor = cnx.cursor()
     cursor.execute("INSERT INTO test (testField) VALUES (123)"
     cursor.close()
@@ -45,29 +46,71 @@
 
 import mysql.connector
 
-from opentelemetry.ext.dbapi import wrap_connect
+from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.ext import dbapi
 from opentelemetry.ext.mysql.version import __version__
 from opentelemetry.trace import TracerProvider, get_tracer
 
 
-def trace_integration(tracer_provider: typing.Optional[TracerProvider] = None):
-    """Integrate with MySQL Connector/Python library.
-       https://dev.mysql.com/doc/connector-python/en/
-    """
-
-    tracer = get_tracer(__name__, __version__, tracer_provider)
-
-    connection_attributes = {
+class MySQLInstrumentor(BaseInstrumentor):
+    _CONNECTION_ATTRIBUTES = {
         "database": "database",
         "port": "server_port",
         "host": "server_host",
         "user": "user",
     }
-    wrap_connect(
-        tracer,
-        mysql.connector,
-        "connect",
-        "mysql",
-        "sql",
-        connection_attributes,
-    )
+
+    _DATABASE_COMPONENT = "mysql"
+    _DATABASE_TYPE = "sql"
+
+    def _instrument(self, **kwargs):
+        """Integrate with MySQL Connector/Python library.
+        https://dev.mysql.com/doc/connector-python/en/
+        """
+        tracer_provider = kwargs.get("tracer_provider")
+
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        dbapi.wrap_connect(
+            tracer,
+            mysql.connector,
+            "connect",
+            self._DATABASE_COMPONENT,
+            self._DATABASE_TYPE,
+            self._CONNECTION_ATTRIBUTES,
+        )
+
+    def _uninstrument(self, **kwargs):
+        """"Disable MySQL instrumentation"""
+        dbapi.unwrap_connect(mysql.connector, "connect")
+
+    # pylint:disable=no-self-use
+    def instrument_connection(self, connection):
+        """Enable instrumentation in a MySQL connection.
+
+        Args:
+            connection: The connection to instrument.
+
+        Returns:
+            An instrumented connection.
+        """
+        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 MySQL connection.
+
+        Args:
+            connection: The connection to uninstrument.
+
+        Returns:
+            An uninstrumented connection.
+        """
+        return dbapi.uninstrument_connection(connection)