fix tests

Author: rahul2393

.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml

@@ -31,3 +31,4 @@ jobs:
           GOOGLE_CLOUD_PROJECT: emulator-test-project
           GOOGLE_CLOUD_TESTS_CREATE_SPANNER_INSTANCE: true
           GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
+          GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true

google/cloud/spanner_v1/database.py

@@ -202,7 +202,6 @@ def __init__(
         self._pool = pool
         pool.bind(self)
 
-        # Initialize session options and sessions manager for multiplexed session support
         self.session_options = SessionOptions()
         self._sessions_manager = DatabaseSessionsManager(self, pool)
 
@@ -767,7 +766,6 @@ def execute_pdml():
             ) as span, MetricsCapture():
                 from google.cloud.spanner_v1.session_options import TransactionType
 
-                # Use sessions manager for partitioned DML operations
                 session = self._sessions_manager.get_session(
                     TransactionType.PARTITIONED
                 )

google/cloud/spanner_v1/session.py

@@ -174,8 +174,13 @@ def create(self):
             request.session.labels = self._labels
 
         observability_options = getattr(self._database, "observability_options", None)
+        span_name = (
+            "CloudSpanner.CreateMultiplexedSession"
+            if self._is_multiplexed
+            else "CloudSpanner.CreateSession"
+        )
         with trace_call(
-            "CloudSpanner.CreateSession",
+            span_name,
             self,
             self._labels,
             observability_options=observability_options,

google/cloud/spanner_v1/session_options.py

@@ -59,10 +59,7 @@ def use_multiplexed(self, transaction_type: TransactionType) -> bool:
         * ENV_VAR_ENABLE_MULTIPLEXED is set to true;
         * ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED is set to true; and
         * multiplexed sessions have not been disabled for partitioned transactions.
-        Multiplexed sessions are enabled for read/write transactions if:
-        * ENV_VAR_ENABLE_MULTIPLEXED is set to true;
-        * ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED is set to true; and
-        * multiplexed sessions have not been disabled for read/write transactions.
+        Multiplexed sessions are **currently disabled** for read / write.
         :type transaction_type: :class:`TransactionType`
         :param transaction_type: the type of transaction to check whether
             multiplexed sessions should be used.
@@ -102,6 +99,10 @@ def disable_multiplexed(
         disable_multiplexed_log_msg_fstring = (
             "Disabling multiplexed sessions for {transaction_type_value} transactions"
         )
+        import logging
+
+        if logger is None:
+            logger = logging.getLogger(__name__)
 
         if transaction_type is None:
             logger.warning(

tests/system/test_observability_options.py

@@ -109,8 +109,23 @@ def test_propagation(enable_extended_tracing):
             len(from_inject_spans) >= 2
         )  # "Expecting at least 2 spans from the injected trace exporter"
         gotNames = [span.name for span in from_inject_spans]
+
+        # Check if multiplexed sessions are enabled
+        import os
+
+        multiplexed_enabled = (
+            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS", "").lower() == "true"
+        )
+
+        # Determine expected session span name based on multiplexed sessions
+        expected_session_span_name = (
+            "CloudSpanner.CreateMultiplexedSession"
+            if multiplexed_enabled
+            else "CloudSpanner.CreateSession"
+        )
+
         wantNames = [
-            "CloudSpanner.CreateSession",
+            expected_session_span_name,
             "CloudSpanner.Snapshot.execute_sql",
         ]
         assert gotNames == wantNames
@@ -392,6 +407,7 @@ def tx_update(txn):
     reason="Tracing requires OpenTelemetry",
 )
 def test_database_partitioned_error():
+    import os
     from opentelemetry.trace.status import StatusCode
 
     db, trace_exporter = create_db_trace_exporter()
@@ -402,43 +418,84 @@ def test_database_partitioned_error():
         pass
 
     got_statuses, got_events = finished_spans_statuses(trace_exporter)
-    # Check for the series of events
-    want_events = [
-        ("Acquiring session", {"kind": "BurstyPool"}),
-        ("Waiting for a session to become available", {"kind": "BurstyPool"}),
-        ("No sessions available in pool. Creating session", {"kind": "BurstyPool"}),
-        ("Creating Session", {}),
-        ("Starting BeginTransaction", {}),
-        (
+
+    multiplexed_partitioned_enabled = (
+        os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS") == "true"
+    )
+
+    if multiplexed_partitioned_enabled:
+        expected_event_names = [
+            "Creating Session",
+            "Using session",
+            "Starting BeginTransaction",
+            "Returning session",
             "exception",
-            {
-                "exception.type": "google.api_core.exceptions.InvalidArgument",
-                "exception.message": "400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
-        (
             "exception",
-            {
-                "exception.type": "google.api_core.exceptions.InvalidArgument",
-                "exception.message": "400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
-    ]
-    assert got_events == want_events
+        ]
+        assert len(got_events) == len(expected_event_names)
+        for i, expected_name in enumerate(expected_event_names):
+            assert got_events[i][0] == expected_name
+
+        assert got_events[1][1]["multiplexed"] is True
+
+        assert got_events[3][1]["multiplexed"] is True
+
+        for i in [4, 5]:
+            assert (
+                got_events[i][1]["exception.type"]
+                == "google.api_core.exceptions.InvalidArgument"
+            )
+            assert (
+                "Table not found: NonExistent" in got_events[i][1]["exception.message"]
+            )
+    else:
+        expected_event_names = [
+            "Acquiring session",
+            "Waiting for a session to become available",
+            "No sessions available in pool. Creating session",
+            "Creating Session",
+            "Using session",
+            "Starting BeginTransaction",
+            "Returning session",
+            "exception",
+            "exception",
+        ]
+
+        assert len(got_events) == len(expected_event_names)
+        for i, expected_name in enumerate(expected_event_names):
+            assert got_events[i][0] == expected_name
+
+        assert got_events[0][1]["kind"] == "BurstyPool"
+        assert got_events[1][1]["kind"] == "BurstyPool"
+        assert got_events[2][1]["kind"] == "BurstyPool"
+
+        assert got_events[4][1]["multiplexed"] is False
+
+        assert got_events[6][1]["multiplexed"] is False
+
+        for i in [7, 8]:
+            assert (
+                got_events[i][1]["exception.type"]
+                == "google.api_core.exceptions.InvalidArgument"
+            )
+            assert (
+                "Table not found: NonExistent" in got_events[i][1]["exception.message"]
+            )
 
-    # Check for the statues.
     codes = StatusCode
+
+    expected_session_span_name = (
+        "CloudSpanner.CreateMultiplexedSession"
+        if multiplexed_partitioned_enabled
+        else "CloudSpanner.CreateSession"
+    )
     want_statuses = [
         (
             "CloudSpanner.Database.execute_partitioned_pdml",
             codes.ERROR,
             "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
         ),
-        ("CloudSpanner.CreateSession", codes.OK, None),
+        (expected_session_span_name, codes.OK, None),
         (
             "CloudSpanner.ExecuteStreamingSql",
             codes.ERROR,