diff --git a/instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py b/instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py index 4e6257fbb1..298ade2b32 100644 --- a/instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py +++ b/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) diff --git a/instrumentation/opentelemetry-instrumentation-aiopg/tests/test_aiopg_integration.py b/instrumentation/opentelemetry-instrumentation-aiopg/tests/test_aiopg_integration.py index f32ef3514b..c5b4b8a154 100644 --- a/instrumentation/opentelemetry-instrumentation-aiopg/tests/test_aiopg_integration.py +++ b/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):