Skip to content

Commit dbd1d23

Browse files
committed
Fix tracing of async cursors for psycopg
This copies the traced_execution of AsyncCursorTracer except query_method is awaited within the span. Fixes open-telemetry#2486
1 parent 81eaea5 commit dbd1d23

File tree

1 file changed

+42
-2
lines changed
  • instrumentation/opentelemetry-instrumentation-psycopg/src/opentelemetry/instrumentation/psycopg

1 file changed

+42
-2
lines changed

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

+42-2
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@
149149
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
150150
from opentelemetry.instrumentation.psycopg.package import _instruments
151151
from opentelemetry.instrumentation.psycopg.version import __version__
152-
from opentelemetry.trace import TracerProvider
152+
from opentelemetry.trace import SpanKind, TracerProvider
153153

154154
_logger = logging.getLogger(__name__)
155155
_OTEL_CURSOR_FACTORY_KEY = "_otel_orig_cursor_factory"
@@ -381,6 +381,46 @@ def callproc(self, *args: Any, **kwargs: Any):
381381
return TracedCursorFactory
382382

383383

384+
class AsyncCursorTracer(CursorTracer):
385+
async def traced_execution(
386+
self,
387+
cursor: dbapi.CursorT,
388+
query_method: Callable[..., Any],
389+
*args: tuple[Any, ...],
390+
**kwargs: dict[Any, Any],
391+
):
392+
name = self.get_operation_name(cursor, args)
393+
if not name:
394+
name = (
395+
self._db_api_integration.database
396+
if self._db_api_integration.database
397+
else self._db_api_integration.name
398+
)
399+
400+
with self._db_api_integration._tracer.start_as_current_span(
401+
name, kind=SpanKind.CLIENT
402+
) as span:
403+
if span.is_recording():
404+
if args and self._commenter_enabled:
405+
if self._enable_attribute_commenter:
406+
# sqlcomment is added to executed query and db.statement span attribute
407+
args = self._update_args_with_added_sql_comment(
408+
args, cursor
409+
)
410+
self._populate_span(span, cursor, *args)
411+
else:
412+
# sqlcomment is only added to executed query
413+
# so db.statement is set before add_sql_comment
414+
self._populate_span(span, cursor, *args)
415+
args = self._update_args_with_added_sql_comment(
416+
args, cursor
417+
)
418+
else:
419+
# no sqlcomment anywhere
420+
self._populate_span(span, cursor, *args)
421+
return await query_method(*args, **kwargs)
422+
423+
384424
def _new_cursor_async_factory(
385425
db_api: DatabaseApiAsyncIntegration | None = None,
386426
base_factory: type[psycopg.AsyncCursor] | None = None,
@@ -395,7 +435,7 @@ def _new_cursor_async_factory(
395435
tracer_provider=tracer_provider,
396436
)
397437
base_factory = base_factory or psycopg.AsyncCursor
398-
_cursor_tracer = CursorTracer(db_api)
438+
_cursor_tracer = AsyncCursorTracer(db_api)
399439

400440
class TracedCursorAsyncFactory(base_factory):
401441
async def execute(self, *args: Any, **kwargs: Any):

0 commit comments

Comments
 (0)