@@ -392,6 +392,7 @@ def tx_update(txn):
392
392
reason = "Tracing requires OpenTelemetry" ,
393
393
)
394
394
def test_database_partitioned_error ():
395
+ import os
395
396
from opentelemetry .trace .status import StatusCode
396
397
397
398
db , trace_exporter = create_db_trace_exporter ()
@@ -402,43 +403,101 @@ def test_database_partitioned_error():
402
403
pass
403
404
404
405
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
- (
406
+
407
+ # Check if multiplexed sessions are enabled for partitioned operations
408
+ multiplexed_partitioned_enabled = (
409
+ os .getenv ("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS" ) == "true"
410
+ )
411
+
412
+ # Define expected events based on whether multiplexed sessions are enabled
413
+ if multiplexed_partitioned_enabled :
414
+ # When multiplexed sessions are enabled for partitioned operations,
415
+ # the execution path is different - sessions manager creates multiplexed sessions directly
416
+ expected_event_names = [
417
+ "Creating Session" ,
418
+ "Using session" ,
419
+ "Starting BeginTransaction" ,
420
+ "Returning session" ,
413
421
"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
422
"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
423
+ ]
424
+ # Check that we have the expected events
425
+ assert len (got_events ) == len (expected_event_names )
426
+ for i , expected_name in enumerate (expected_event_names ):
427
+ assert got_events [i ][0 ] == expected_name
428
+
429
+ # Verify session usage event shows multiplexed session
430
+ assert got_events [1 ][1 ]["multiplexed" ] is True
431
+
432
+ # Verify session return event shows multiplexed session
433
+ assert got_events [3 ][1 ]["multiplexed" ] is True
434
+
435
+ # Verify the exception details
436
+ for i in [4 , 5 ]: # Both exception events
437
+ assert (
438
+ got_events [i ][1 ]["exception.type" ]
439
+ == "google.api_core.exceptions.InvalidArgument"
440
+ )
441
+ assert (
442
+ "Table not found: NonExistent" in got_events [i ][1 ]["exception.message" ]
443
+ )
444
+ else :
445
+ # When multiplexed sessions are disabled, sessions manager still manages sessions
446
+ # but uses regular pool sessions instead of multiplexed sessions
447
+ expected_event_names = [
448
+ "Acquiring session" ,
449
+ "Waiting for a session to become available" ,
450
+ "No sessions available in pool. Creating session" ,
451
+ "Creating Session" ,
452
+ "Using session" ,
453
+ "Starting BeginTransaction" ,
454
+ "Returning session" ,
455
+ "exception" ,
456
+ "exception" ,
457
+ ]
432
458
433
- # Check for the statues.
459
+ # Check that we have the expected events
460
+ assert len (got_events ) == len (expected_event_names )
461
+ for i , expected_name in enumerate (expected_event_names ):
462
+ assert got_events [i ][0 ] == expected_name
463
+
464
+ # Verify pool-related events
465
+ assert got_events [0 ][1 ]["kind" ] == "BurstyPool"
466
+ assert got_events [1 ][1 ]["kind" ] == "BurstyPool"
467
+ assert got_events [2 ][1 ]["kind" ] == "BurstyPool"
468
+
469
+ # Verify session usage event shows non-multiplexed session
470
+ assert got_events [4 ][1 ]["multiplexed" ] is False
471
+
472
+ # Verify session return event shows non-multiplexed session
473
+ assert got_events [6 ][1 ]["multiplexed" ] is False
474
+
475
+ # Verify the exception details
476
+ for i in [7 , 8 ]: # Both exception events
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
+ )
484
+
485
+ # Check for the statuses.
434
486
codes = StatusCode
487
+
488
+ # Determine expected session creation span name based on multiplexed sessions
489
+ expected_session_span_name = (
490
+ "CloudSpanner.CreateMultiplexedSession"
491
+ if multiplexed_partitioned_enabled
492
+ else "CloudSpanner.CreateSession"
493
+ )
435
494
want_statuses = [
436
495
(
437
496
"CloudSpanner.Database.execute_partitioned_pdml" ,
438
497
codes .ERROR ,
439
498
"InvalidArgument: 400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
440
499
),
441
- ("CloudSpanner.CreateSession" , codes .OK , None ),
500
+ (expected_session_span_name , codes .OK , None ),
442
501
(
443
502
"CloudSpanner.ExecuteStreamingSql" ,
444
503
codes .ERROR ,
0 commit comments