Refactor TracedConnectionProxy

Fixes #1077

Author: ocelotl

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

@@ -214,8 +214,8 @@ def uninstrument_connection(connection):
     Returns:
         An uninstrumented connection.
     """
-    if isinstance(connection, wrapt.ObjectProxy):
-        return connection.__wrapped__
+    if isinstance(connection, _TracedConnectionProxy):
+        return connection._connection
 
     _logger.warning("Connection is not instrumented")
     return connection
@@ -300,28 +300,36 @@ def get_connection_attributes(self, connection):
             self.span_attributes[SpanAttributes.NET_PEER_PORT] = port
 
 
+class _TracedConnectionProxy:
+    pass
+
+
 def get_traced_connection_proxy(
     connection, db_api_integration, *args, **kwargs
 ):
     # pylint: disable=abstract-method
-    class TracedConnectionProxy(wrapt.ObjectProxy):
-        # pylint: disable=unused-argument
-        def __init__(self, connection, *args, **kwargs):
-            wrapt.ObjectProxy.__init__(self, connection)
+    class TracedConnectionProxy(type(connection), _TracedConnectionProxy):
+        def __init__(self, connection):
+            self._connection = connection
 
-        def cursor(self, *args, **kwargs):
-            return get_traced_cursor_proxy(
-                self.__wrapped__.cursor(*args, **kwargs), db_api_integration
+        def __getattr__(self, name):
+            return object.__getattribute__(
+                object.__getattribute__(self, "_connection"), name
             )
 
         def __enter__(self):
-            self.__wrapped__.__enter__()
-            return self
+            self._connection.__enter__()
+            return self._connection
 
-        def __exit__(self, *args, **kwargs):
-            self.__wrapped__.__exit__(*args, **kwargs)
+        def __exit__(self, exc_type, exc_value, traceback):
+            self._connection.__exit__(exc_type, exc_value, traceback)
+
+        def cursor(self, *args, **kwargs):
+            return get_traced_cursor_proxy(
+                self._connection.cursor(*args, **kwargs), db_api_integration
+            )
 
-    return TracedConnectionProxy(connection, *args, **kwargs)
+    return TracedConnectionProxy(connection)
 
 
 class CursorTracer:

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

@@ -262,14 +262,14 @@ def test_callproc(self):
 
     @mock.patch("opentelemetry.instrumentation.dbapi")
     def test_wrap_connect(self, mock_dbapi):
-        dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
+        dbapi.wrap_connect(self.tracer, MockConnectionEmpty(), "connect", "-")
         connection = mock_dbapi.connect()
         self.assertEqual(mock_dbapi.connect.call_count, 1)
-        self.assertIsInstance(connection.__wrapped__, mock.Mock)
+        self.assertIsInstance(connection._connection, mock.Mock)
 
     @mock.patch("opentelemetry.instrumentation.dbapi")
     def test_unwrap_connect(self, mock_dbapi):
-        dbapi.wrap_connect(self.tracer, mock_dbapi, "connect", "-")
+        dbapi.wrap_connect(self.tracer, MockConnectionEmpty(), "connect", "-")
         connection = mock_dbapi.connect()
         self.assertEqual(mock_dbapi.connect.call_count, 1)
 
@@ -279,19 +279,19 @@ def test_unwrap_connect(self, mock_dbapi):
         self.assertIsInstance(connection, mock.Mock)
 
     def test_instrument_connection(self):
-        connection = mock.Mock()
+        connection = MockConnectionEmpty()
         # Avoid get_attributes failing because can't concatenate mock
         connection.database = "-"
         connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
-        self.assertIs(connection2.__wrapped__, connection)
+        self.assertIs(connection2._connection, connection)
 
     def test_uninstrument_connection(self):
-        connection = mock.Mock()
+        connection = MockConnectionEmpty()
         # Set connection.database to avoid a failure because mock can't
         # be concatenated
         connection.database = "-"
         connection2 = dbapi.instrument_connection(self.tracer, connection, "-")
-        self.assertIs(connection2.__wrapped__, connection)
+        self.assertIs(connection2._connection, connection)
 
         connection3 = dbapi.uninstrument_connection(connection2)
         self.assertIs(connection3, connection)
@@ -307,10 +307,12 @@ def mock_connect(*args, **kwargs):
     server_host = kwargs.get("server_host")
     server_port = kwargs.get("server_port")
     user = kwargs.get("user")
-    return MockConnection(database, server_port, server_host, user)
+    return MockConnectionWithAttributes(
+        database, server_port, server_host, user
+    )
 
 
-class MockConnection:
+class MockConnectionWithAttributes:
     def __init__(self, database, server_port, server_host, user):
         self.database = database
         self.server_port = server_port
@@ -343,3 +345,7 @@ def executemany(self, query, params=None, throw_exception=False):
     def callproc(self, query, params=None, throw_exception=False):
         if throw_exception:
             raise Exception("Test Exception")
+
+
+class MockConnectionEmpty:
+    pass