149149from opentelemetry .instrumentation .instrumentor import BaseInstrumentor
150150from opentelemetry .instrumentation .psycopg .package import _instruments
151151from 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+
384424def _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