149
149
from opentelemetry .instrumentation .instrumentor import BaseInstrumentor
150
150
from opentelemetry .instrumentation .psycopg .package import _instruments
151
151
from opentelemetry .instrumentation .psycopg .version import __version__
152
- from opentelemetry .trace import TracerProvider
152
+ from opentelemetry .trace import SpanKind , TracerProvider
153
153
154
154
_logger = logging .getLogger (__name__ )
155
155
_OTEL_CURSOR_FACTORY_KEY = "_otel_orig_cursor_factory"
@@ -381,6 +381,46 @@ def callproc(self, *args: Any, **kwargs: Any):
381
381
return TracedCursorFactory
382
382
383
383
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
+
384
424
def _new_cursor_async_factory (
385
425
db_api : DatabaseApiAsyncIntegration | None = None ,
386
426
base_factory : type [psycopg .AsyncCursor ] | None = None ,
@@ -395,7 +435,7 @@ def _new_cursor_async_factory(
395
435
tracer_provider = tracer_provider ,
396
436
)
397
437
base_factory = base_factory or psycopg .AsyncCursor
398
- _cursor_tracer = CursorTracer (db_api )
438
+ _cursor_tracer = AsyncCursorTracer (db_api )
399
439
400
440
class TracedCursorAsyncFactory (base_factory ):
401
441
async def execute (self , * args : Any , ** kwargs : Any ):
0 commit comments