Populate db.statement when Composable passed to cursor.execute

Author: Oliver Paraskos

instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py

@@ -3,6 +3,7 @@
 from collections.abc import Coroutine
 
 import wrapt
+from psycopg2.sql import Composable  # pylint: disable=no-name-in-module
 
 from opentelemetry.instrumentation.dbapi import (
     CursorTracer,
@@ -121,6 +122,18 @@ async def traced_execution(
             self._populate_span(span, cursor, *args)
             return await query_method(*args, **kwargs)
 
+    def get_operation_name(self, cursor, args):
+        if len(args) and isinstance(args[0], Composable):
+            return args[0].as_string(cursor)
+
+        return super().get_operation_name(cursor, args)
+
+    def get_statement(self, cursor, args):
+        if len(args) and isinstance(args[0], Composable):
+            return args[0].as_string(cursor)
+
+        return super().get_statement(cursor, args)
+
 
 def get_traced_cursor_proxy(cursor, db_api_integration, *args, **kwargs):
     _traced_cursor = AsyncCursorTracer(db_api_integration)

instrumentation/opentelemetry-instrumentation-aiopg/tests/test_aiopg_integration.py

@@ -495,6 +495,24 @@ def test_uninstrument_connection(self):
             connection4 = wrappers.uninstrument_connection(connection)
         self.assertIs(connection4, connection)
 
+    def test_composable(self):
+        from psycopg2 import sql
+
+        db_integration = AiopgIntegration(self.tracer, "testcomponent")
+        mock_connection = async_call(
+            db_integration.wrapped_connection(mock_connect, {}, {})
+        )
+        cursor = async_call(mock_connection.cursor())
+        async_call(cursor.execute(sql.SQL("SELECT * FROM my_table")))
+        spans_list = self.memory_exporter.get_finished_spans()
+
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            "SELECT * FROM my_table",
+        )
+
 
 # pylint: disable=unused-argument
 async def mock_connect(*args, **kwargs):