Fix tracing of async cursors for psycopg (#3324)

* Fix tracing of async cursors for psycopg

This copies the traced_execution of AsyncCursorTracer except
query_method is awaited within the span.

Fixes #2486

* Adding a regression test for asynccursor bug

Adds a test to check that the async tracer is actually awaiting
cursor.execute.

On my machine, with the bug present, the duration is 14000ns. With the
bug patched the the duration is orders of magnitude larger.

* Make return type of query_method an awaitable

* Move implementation to traced_execution_async

---------

Co-authored-by: Aaron Abbott <[email protected]>

Author: ibash

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

@@ -42,7 +42,7 @@
 import functools
 import logging
 import re
-from typing import Any, Callable, Generic, TypeVar
+from typing import Any, Awaitable, Callable, Generic, TypeVar
 
 import wrapt
 from wrapt import wrap_function_wrapper
@@ -596,6 +596,44 @@ def traced_execution(
                     self._populate_span(span, cursor, *args)
             return query_method(*args, **kwargs)
 
+    async def traced_execution_async(
+        self,
+        cursor: CursorT,
+        query_method: Callable[..., Awaitable[Any]],
+        *args: tuple[Any, ...],
+        **kwargs: dict[Any, Any],
+    ):
+        name = self.get_operation_name(cursor, args)
+        if not name:
+            name = (
+                self._db_api_integration.database
+                if self._db_api_integration.database
+                else self._db_api_integration.name
+            )
+
+        with self._db_api_integration._tracer.start_as_current_span(
+            name, kind=SpanKind.CLIENT
+        ) as span:
+            if span.is_recording():
+                if args and self._commenter_enabled:
+                    if self._enable_attribute_commenter:
+                        # sqlcomment is added to executed query and db.statement span attribute
+                        args = self._update_args_with_added_sql_comment(
+                            args, cursor
+                        )
+                        self._populate_span(span, cursor, *args)
+                    else:
+                        # sqlcomment is only added to executed query
+                        # so db.statement is set before add_sql_comment
+                        self._populate_span(span, cursor, *args)
+                        args = self._update_args_with_added_sql_comment(
+                            args, cursor
+                        )
+                else:
+                    # no sqlcomment anywhere
+                    self._populate_span(span, cursor, *args)
+            return await query_method(*args, **kwargs)
+
 
 # pylint: disable=abstract-method
 class TracedCursorProxy(wrapt.ObjectProxy, Generic[CursorT]):

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

@@ -399,17 +399,17 @@ def _new_cursor_async_factory(
 
     class TracedCursorAsyncFactory(base_factory):
         async def execute(self, *args: Any, **kwargs: Any):
-            return await _cursor_tracer.traced_execution(
+            return await _cursor_tracer.traced_execution_async(
                 self, super().execute, *args, **kwargs
             )
 
         async def executemany(self, *args: Any, **kwargs: Any):
-            return await _cursor_tracer.traced_execution(
+            return await _cursor_tracer.traced_execution_async(
                 self, super().executemany, *args, **kwargs
             )
 
         async def callproc(self, *args: Any, **kwargs: Any):
-            return await _cursor_tracer.traced_execution(
+            return await _cursor_tracer.traced_execution_async(
                 self, super().callproc, *args, **kwargs
             )
 

instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import types
 from unittest import IsolatedAsyncioTestCase, mock
 
@@ -50,10 +51,15 @@ def __init__(self, *args, **kwargs):
         pass
 
     # pylint: disable=unused-argument, no-self-use
-    async def execute(self, query, params=None, throw_exception=False):
+    async def execute(
+        self, query, params=None, throw_exception=False, delay=0.0
+    ):
         if throw_exception:
             raise psycopg.Error("Test Exception")
 
+        if delay:
+            await asyncio.sleep(delay)
+
     # pylint: disable=unused-argument, no-self-use
     async def executemany(self, query, params=None, throw_exception=False):
         if throw_exception:
@@ -492,3 +498,27 @@ async def test_not_recording_async(self):
             self.assertFalse(mock_span.set_status.called)
 
         PsycopgInstrumentor().uninstrument()
+
+    async def test_tracing_is_async(self):
+        PsycopgInstrumentor().instrument()
+
+        # before this async fix cursor.execute would take 14000 ns, delaying for
+        # 100,000ns
+        delay = 0.0001
+
+        async def test_async_connection():
+            acnx = await psycopg.AsyncConnection.connect("test")
+            async with acnx as cnx:
+                async with cnx.cursor() as cursor:
+                    await cursor.execute("SELECT * FROM test", delay=delay)
+
+        await test_async_connection()
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 1)
+        span = spans_list[0]
+
+        # duration is nanoseconds
+        duration = span.end_time - span.start_time
+        self.assertGreater(duration, delay * 1e9)
+
+        PsycopgInstrumentor().uninstrument()