@@ -109,8 +109,23 @@ def test_propagation(enable_extended_tracing):
109
109
len (from_inject_spans ) >= 2
110
110
) # "Expecting at least 2 spans from the injected trace exporter"
111
111
gotNames = [span .name for span in from_inject_spans ]
112
+
113
+ # Check if multiplexed sessions are enabled
114
+ import os
115
+
116
+ multiplexed_enabled = (
117
+ os .getenv ("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS" , "" ).lower () == "true"
118
+ )
119
+
120
+ # Determine expected session span name based on multiplexed sessions
121
+ expected_session_span_name = (
122
+ "CloudSpanner.CreateMultiplexedSession"
123
+ if multiplexed_enabled
124
+ else "CloudSpanner.CreateSession"
125
+ )
126
+
112
127
wantNames = [
113
- "CloudSpanner.CreateSession" ,
128
+ expected_session_span_name ,
114
129
"CloudSpanner.Snapshot.execute_sql" ,
115
130
]
116
131
assert gotNames == wantNames
@@ -392,6 +407,7 @@ def tx_update(txn):
392
407
reason = "Tracing requires OpenTelemetry" ,
393
408
)
394
409
def test_database_partitioned_error ():
410
+ import os
395
411
from opentelemetry .trace .status import StatusCode
396
412
397
413
db , trace_exporter = create_db_trace_exporter ()
@@ -402,43 +418,84 @@ def test_database_partitioned_error():
402
418
pass
403
419
404
420
got_statuses , got_events = finished_spans_statuses (trace_exporter )
405
- # Check for the series of events
406
- want_events = [
407
- ("Acquiring session" , {"kind" : "BurstyPool" }),
408
- ("Waiting for a session to become available" , {"kind" : "BurstyPool" }),
409
- ("No sessions available in pool. Creating session" , {"kind" : "BurstyPool" }),
410
- ("Creating Session" , {}),
411
- ("Starting BeginTransaction" , {}),
412
- (
421
+
422
+ multiplexed_partitioned_enabled = (
423
+ os .getenv ("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS" ) == "true"
424
+ )
425
+
426
+ if multiplexed_partitioned_enabled :
427
+ expected_event_names = [
428
+ "Creating Session" ,
429
+ "Using session" ,
430
+ "Starting BeginTransaction" ,
431
+ "Returning session" ,
413
432
"exception" ,
414
- {
415
- "exception.type" : "google.api_core.exceptions.InvalidArgument" ,
416
- "exception.message" : "400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
417
- "exception.stacktrace" : "EPHEMERAL" ,
418
- "exception.escaped" : "False" ,
419
- },
420
- ),
421
- (
422
433
"exception" ,
423
- {
424
- "exception.type" : "google.api_core.exceptions.InvalidArgument" ,
425
- "exception.message" : "400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
426
- "exception.stacktrace" : "EPHEMERAL" ,
427
- "exception.escaped" : "False" ,
428
- },
429
- ),
430
- ]
431
- assert got_events == want_events
434
+ ]
435
+ assert len (got_events ) == len (expected_event_names )
436
+ for i , expected_name in enumerate (expected_event_names ):
437
+ assert got_events [i ][0 ] == expected_name
438
+
439
+ assert got_events [1 ][1 ]["multiplexed" ] is True
440
+
441
+ assert got_events [3 ][1 ]["multiplexed" ] is True
442
+
443
+ for i in [4 , 5 ]:
444
+ assert (
445
+ got_events [i ][1 ]["exception.type" ]
446
+ == "google.api_core.exceptions.InvalidArgument"
447
+ )
448
+ assert (
449
+ "Table not found: NonExistent" in got_events [i ][1 ]["exception.message" ]
450
+ )
451
+ else :
452
+ expected_event_names = [
453
+ "Acquiring session" ,
454
+ "Waiting for a session to become available" ,
455
+ "No sessions available in pool. Creating session" ,
456
+ "Creating Session" ,
457
+ "Using session" ,
458
+ "Starting BeginTransaction" ,
459
+ "Returning session" ,
460
+ "exception" ,
461
+ "exception" ,
462
+ ]
463
+
464
+ assert len (got_events ) == len (expected_event_names )
465
+ for i , expected_name in enumerate (expected_event_names ):
466
+ assert got_events [i ][0 ] == expected_name
467
+
468
+ assert got_events [0 ][1 ]["kind" ] == "BurstyPool"
469
+ assert got_events [1 ][1 ]["kind" ] == "BurstyPool"
470
+ assert got_events [2 ][1 ]["kind" ] == "BurstyPool"
471
+
472
+ assert got_events [4 ][1 ]["multiplexed" ] is False
473
+
474
+ assert got_events [6 ][1 ]["multiplexed" ] is False
475
+
476
+ for i in [7 , 8 ]:
477
+ assert (
478
+ got_events [i ][1 ]["exception.type" ]
479
+ == "google.api_core.exceptions.InvalidArgument"
480
+ )
481
+ assert (
482
+ "Table not found: NonExistent" in got_events [i ][1 ]["exception.message" ]
483
+ )
432
484
433
- # Check for the statues.
434
485
codes = StatusCode
486
+
487
+ expected_session_span_name = (
488
+ "CloudSpanner.CreateMultiplexedSession"
489
+ if multiplexed_partitioned_enabled
490
+ else "CloudSpanner.CreateSession"
491
+ )
435
492
want_statuses = [
436
493
(
437
494
"CloudSpanner.Database.execute_partitioned_pdml" ,
438
495
codes .ERROR ,
439
496
"InvalidArgument: 400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
440
497
),
441
- ("CloudSpanner.CreateSession" , codes .OK , None ),
498
+ (expected_session_span_name , codes .OK , None ),
442
499
(
443
500
"CloudSpanner.ExecuteStreamingSql" ,
444
501
codes .ERROR ,
0 commit comments