feat(asyncpg): add capture_connection_cleanup option to silence pool cleanup spans

Closes #4264

Author: RiyaChaturvedi37

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `opentelemetry-instrumentation-asyncpg`: Add `capture_connection_cleanup` option to silence connection cleanup spans when using a pool
+  ([#4264](https://github.com/open-telemetry/opentelemetry-python-contrib/issues/4264))
 - `opentelemetry-instrumentation-confluent-kafka`: Loosen confluent-kafka upper bound to <3.0.0
   ([#4289](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4289))
 - `opentelemetry-instrumentation`: Add support for wrapt 2.x

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

@@ -51,10 +51,10 @@ async def main():
 
 import re
 from typing import Collection
-
+ 
 import asyncpg
 import wrapt
-
+ 
 from opentelemetry import trace
 from opentelemetry.instrumentation.asyncpg.package import _instruments
 from opentelemetry.instrumentation.asyncpg.version import __version__
@@ -75,23 +75,23 @@ async def main():
 )
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
-
+ 
+ 
 def _hydrate_span_from_args(connection, query, parameters) -> dict:
     """Get network and database attributes from connection."""
     span_attributes = {DB_SYSTEM: DbSystemValues.POSTGRESQL.value}
-
+ 
     # connection contains _params attribute which is a namedtuple ConnectionParameters.
     # https://github.com/MagicStack/asyncpg/blob/master/asyncpg/connection.py#L68
-
+ 
     params = getattr(connection, "_params", None)
     dbname = getattr(params, "database", None)
     if dbname:
         span_attributes[DB_NAME] = dbname
     user = getattr(params, "user", None)
     if user:
         span_attributes[DB_USER] = user
-
+ 
     # connection contains _addr attribute which is either a host/port tuple, or unix socket string
     # https://magicstack.github.io/asyncpg/current/_modules/asyncpg/connection.html
     addr = getattr(connection, "_addr", None)
@@ -102,27 +102,39 @@ def _hydrate_span_from_args(connection, query, parameters) -> dict:
     elif isinstance(addr, str):
         span_attributes[NET_PEER_NAME] = addr
         span_attributes[NET_TRANSPORT] = NetTransportValues.OTHER.value
-
+ 
     if query is not None:
         span_attributes[DB_STATEMENT] = query
-
+ 
     if parameters is not None and len(parameters) > 0:
         span_attributes["db.statement.parameters"] = str(parameters)
-
+ 
     return span_attributes
-
-
+ 
+ 
 class AsyncPGInstrumentor(BaseInstrumentor):
     _leading_comment_remover = re.compile(r"^/\*.*?\*/")
+    _CLEANUP_QUERIES = frozenset([
+        "SELECT pg_advisory_unlock_all()",
+        "CLOSE ALL",
+        "UNLISTEN *",
+        "RESET ALL",
+    ])
     _tracer = None
-
-    def __init__(self, capture_parameters=False):
+ 
+    def _is_cleanup_query(self, query: str) -> bool:
+        if query is None:
+            return False
+        return any(q in query for q in self._CLEANUP_QUERIES)
+ 
+    def __init__(self, capture_parameters=False, capture_connection_cleanup=True):
         super().__init__()
         self.capture_parameters = capture_parameters
-
+        self.capture_connection_cleanup = capture_connection_cleanup
+ 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
-
+ 
     def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")
         self._tracer = trace.get_tracer(
@@ -131,7 +143,7 @@ def _instrument(self, **kwargs):
             tracer_provider,
             schema_url="https://opentelemetry.io/schemas/1.11.0",
         )
-
+ 
         for method in [
             "Connection.execute",
             "Connection.executemany",
@@ -142,7 +154,7 @@ def _instrument(self, **kwargs):
             wrapt.wrap_function_wrapper(
                 "asyncpg.connection", method, self._do_execute
             )
-
+ 
         for method in [
             "Cursor.fetch",
             "Cursor.forward",
@@ -152,7 +164,7 @@ def _instrument(self, **kwargs):
             wrapt.wrap_function_wrapper(
                 "asyncpg.cursor", method, self._do_cursor_execute
             )
-
+ 
     def _uninstrument(self, **__):
         for cls, methods in [
             (
@@ -164,27 +176,30 @@ def _uninstrument(self, **__):
         ]:
             for method_name in methods:
                 unwrap(cls, method_name)
-
+ 
     async def _do_execute(self, func, instance, args, kwargs):
         exception = None
         params = getattr(instance, "_params", None)
         name = (
             args[0] if args[0] else getattr(params, "database", "postgresql")
         )
-
+ 
         try:
             # Strip leading comments so we get the operation name.
             name = self._leading_comment_remover.sub("", name).split()[0]
         except IndexError:
             name = ""
-
+ 
         # Hydrate attributes before span creation to enable filtering
         span_attributes = _hydrate_span_from_args(
             instance,
             args[0],
             args[1:] if self.capture_parameters else None,
         )
-
+ 
+        if not self.capture_connection_cleanup and self._is_cleanup_query(args[0]):
+            return await func(*args, **kwargs)
+ 
         with self._tracer.start_as_current_span(
             name, kind=SpanKind.CLIENT, attributes=span_attributes
         ) as span:
@@ -196,9 +211,9 @@ async def _do_execute(self, func, instance, args, kwargs):
             finally:
                 if span.is_recording() and exception is not None:
                     span.set_status(Status(StatusCode.ERROR))
-
+ 
         return result
-
+ 
     async def _do_cursor_execute(self, func, instance, args, kwargs):
         """Wrap cursor based functions. For every call this will generate a new span."""
         exception = None
@@ -208,20 +223,20 @@ async def _do_cursor_execute(self, func, instance, args, kwargs):
             if instance._query
             else getattr(params, "database", "postgresql")
         )
-
+ 
         try:
             # Strip leading comments so we get the operation name.
             name = self._leading_comment_remover.sub("", name).split()[0]
         except IndexError:
             name = ""
-
+ 
         # Hydrate attributes before span creation to enable filtering
         span_attributes = _hydrate_span_from_args(
             instance._connection,
             instance._query,
             instance._args if self.capture_parameters else None,
         )
-
+ 
         stop = False
         with self._tracer.start_as_current_span(
             f"CURSOR: {name}",
@@ -239,7 +254,8 @@ async def _do_cursor_execute(self, func, instance, args, kwargs):
             finally:
                 if span.is_recording() and exception is not None:
                     span.set_status(Status(StatusCode.ERROR))
-
+ 
         if not stop:
             return result
         raise StopAsyncIteration
+    

instrumentation/opentelemetry-instrumentation-asyncpg/tests/test_asyncpg_wrapper.py

@@ -143,4 +143,31 @@ async def exec_mock(*args, **kwargs):
         asyncio.run(crs.fetch(1))
 
         spans = self.memory_exporter.get_finished_spans()
+        def test_capture_connection_cleanup_false(self):
+        """Test that cleanup queries are not traced when capture_connection_cleanup=False."""
+        AsyncPGInstrumentor().uninstrument()
+        apg = AsyncPGInstrumentor(capture_connection_cleanup=False)
+        apg.instrument(tracer_provider=self.tracer_provider)
+
+        async def mock_execute(*args, **kwargs):
+            return None
+
+        conn = mock.Mock()
+        conn._params = mock.Mock()
+        conn._params.database = "testdb"
+        conn._params.user = "testuser"
+        conn._addr = ("localhost", 5432)
+
+        for cleanup_query in [
+            "SELECT pg_advisory_unlock_all()",
+            "CLOSE ALL",
+            "UNLISTEN *",
+            "RESET ALL",
+        ]:
+            asyncio.run(
+                apg._do_execute(mock_execute, conn, (cleanup_query,), {})
+            )
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
         self.assertEqual(len(spans), 0)