Refactor abstract classes BaseTracedConnection|CursorProxy

Author: tammy-baylis-swi

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

@@ -41,6 +41,7 @@
 import logging
 import re
 import typing
+from abc import ABC, abstractmethod
 
 import wrapt
 
@@ -390,35 +391,41 @@ def get_connection_attributes(self, connection):
             self.span_attributes[SpanAttributes.NET_PEER_PORT] = port
 
 
+class BaseTracedConnectionProxy(ABC, wrapt.ObjectProxy):
+    # pylint: disable=unused-argument
+    def __init__(self, connection, *args, **kwargs):
+        wrapt.ObjectProxy.__init__(self, connection)
+
+    def __getattribute__(self, name):
+        if object.__getattribute__(self, name):
+            return object.__getattribute__(self, name)
+
+        return object.__getattribute__(
+            object.__getattribute__(self, "_connection"), name
+        )
+
+    @abstractmethod
+    def cursor(self, *args, **kwargs):
+        """Returns instrumented database query cursor"""
+
+    def __enter__(self):
+        self.__wrapped__.__enter__()
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.__wrapped__.__exit__(*args, **kwargs)
+
+
 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)
-
-        def __getattribute__(self, name):
-            if object.__getattribute__(self, name):
-                return object.__getattribute__(self, name)
-
-            return object.__getattribute__(
-                object.__getattribute__(self, "_connection"), name
-            )
-
+    class TracedConnectionProxy(BaseTracedConnectionProxy):
         def cursor(self, *args, **kwargs):
             return get_traced_cursor_proxy(
                 self.__wrapped__.cursor(*args, **kwargs), db_api_integration
             )
 
-        def __enter__(self):
-            self.__wrapped__.__enter__()
-            return self
-
-        def __exit__(self, *args, **kwargs):
-            self.__wrapped__.__exit__(*args, **kwargs)
-
     return TracedConnectionProxy(connection, *args, **kwargs)
 
 
@@ -543,35 +550,49 @@ def traced_execution(
             return query_method(*args, **kwargs)
 
 
-def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
-    _cursor_tracer = CursorTracer(db_api_integration)
+# pylint: disable=abstract-method
+class BaseTracedCursorProxy(ABC, wrapt.ObjectProxy):
+    @abstractmethod
+    # def _set_cursor_tracer(self, db_api_integration, **kwargs):
+    def _set_cursor_tracer(self, **kwargs):
+        """Set self._cursor_tracer with db_api_integration kwarg"""
 
-    # pylint: disable=abstract-method
-    class TracedCursorProxy(wrapt.ObjectProxy):
-        # pylint: disable=unused-argument
-        def __init__(self, cursor, *args, **kwargs):
-            wrapt.ObjectProxy.__init__(self, cursor)
-
-        def execute(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.execute, *args, **kwargs
-            )
+    # pylint: disable=unused-argument
+    def __init__(self, cursor, *args, **kwargs):
+        wrapt.ObjectProxy.__init__(self, cursor)
+        self._cursor_tracer = self._set_cursor_tracer(**kwargs)
 
-        def executemany(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.executemany, *args, **kwargs
-            )
+    def callproc(self, *args, **kwargs):
+        return self._cursor_tracer.traced_execution(
+            self.__wrapped__, self.__wrapped__.callproc, *args, **kwargs
+        )
 
-        def callproc(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.callproc, *args, **kwargs
-            )
+    def execute(self, *args, **kwargs):
+        return self._cursor_tracer.traced_execution(
+            self.__wrapped__, self.__wrapped__.execute, *args, **kwargs
+        )
+
+    def executemany(self, *args, **kwargs):
+        return self._cursor_tracer.traced_execution(
+            self.__wrapped__, self.__wrapped__.executemany, *args, **kwargs
+        )
+
+    def __enter__(self):
+        self.__wrapped__.__enter__()
+        return self
 
-        def __enter__(self):
-            self.__wrapped__.__enter__()
-            return self
+    def __exit__(self, *args, **kwargs):
+        self.__wrapped__.__exit__(*args, **kwargs)
 
-        def __exit__(self, *args, **kwargs):
-            self.__wrapped__.__exit__(*args, **kwargs)
 
-    return TracedCursorProxy(cursor, *args, **kwargs)
+def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
+    class TracedCursorProxy(BaseTracedCursorProxy):
+        def _set_cursor_tracer(self, **kwargs):
+            self._cursor_tracer = CursorTracer(
+                kwargs.get("db_api_integration"),
+            )
+
+    kwargs["db_api_integration"] = db_api_integration
+    cursor_proxy = TracedCursorProxy(cursor, *args, **kwargs)
+    cursor_proxy._set_cursor_tracer(**kwargs)
+    return cursor_proxy

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

@@ -231,19 +231,7 @@ 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)
-
-        def __getattribute__(self, name):
-            if object.__getattribute__(self, name):
-                return object.__getattribute__(self, name)
-
-            return object.__getattribute__(
-                object.__getattribute__(self, "_connection"), name
-            )
-
+    class TracedConnectionProxy(dbapi.BaseTracedConnectionProxy):
         def cursor(self, *args, **kwargs):
             wrapped_cursor = self.__wrapped__.cursor(*args, **kwargs)
 
@@ -306,13 +294,6 @@ def is_mysql_connector_cursor_prepared(self, cursor):  # pylint: disable=no-self
 
             return False
 
-        def __enter__(self):
-            self.__wrapped__.__enter__()
-            return self
-
-        def __exit__(self, *args, **kwargs):
-            self.__wrapped__.__exit__(*args, **kwargs)
-
     return TracedConnectionProxy(connection, *args, **kwargs)
 
 
@@ -325,40 +306,19 @@ def __init__(
         super().__init__(db_api_integration)
         # It's common to have multiple db client cursors per app,
         # so enable_commenter is set at the cursor level and used
-        # during traced query execution for mysql-connector
+        # during CursorTracer.traced_execution for mysql-connector
         self._commenter_enabled = enable_commenter
 
 
 def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
-    enable_commenter = kwargs.get("enable_commenter", False)
-    _cursor_tracer = CursorTracer(db_api_integration, enable_commenter)
-
-    # pylint: disable=abstract-method
-    class TracedCursorProxy(wrapt.ObjectProxy):
-        # pylint: disable=unused-argument
-        def __init__(self, cursor, *args, **kwargs):
-            wrapt.ObjectProxy.__init__(self, cursor)
-
-        def execute(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.execute, *args, **kwargs
-            )
-
-        def executemany(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.executemany, *args, **kwargs
+    class TracedCursorProxy(dbapi.BaseTracedCursorProxy):
+        def _set_cursor_tracer(self, **kwargs):
+            self._cursor_tracer = CursorTracer(
+                kwargs.get("db_api_integration"),
+                kwargs.get("enable_commenter", False),
             )
 
-        def callproc(self, *args, **kwargs):
-            return _cursor_tracer.traced_execution(
-                self.__wrapped__, self.__wrapped__.callproc, *args, **kwargs
-            )
-
-        def __enter__(self):
-            self.__wrapped__.__enter__()
-            return self
-
-        def __exit__(self, *args, **kwargs):
-            self.__wrapped__.__exit__(*args, **kwargs)
-
-    return TracedCursorProxy(cursor, *args, **kwargs)
+    kwargs["db_api_integration"] = db_api_integration
+    cursor_proxy = TracedCursorProxy(cursor, *args, **kwargs)
+    cursor_proxy._set_cursor_tracer(**kwargs)
+    return cursor_proxy