ext/{dbapi,pymysql}: Implement methods to (un)-instrument connections

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

@@ -2,5 +2,6 @@
 
 ## Unreleased
 
+- Implement instrument_connection and uninstrument_connection ([#624](https://github.com/open-telemetry/opentelemetry-python/pull/624))
 - 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/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,46 @@ 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.
+
+        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 PyMySQL connection.
+
+        Args:
+            connection: The connection to uninstrument.
+
+        Returns:
+            An uninstrumented connection.
+        """
+        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)