more cleanup

Author: rahul2393

tests/mockserver_tests/test_basics.py

@@ -36,7 +36,6 @@
     unavailable_status,
     add_execute_streaming_sql_results,
 )
-from tests._helpers import is_multiplexed_enabled
 
 
 class TestBasics(MockServerTestBase):
@@ -105,11 +104,8 @@ def test_dbapi_partitioned_dml(self):
             requests[idx], BatchCreateSessionsRequest
         ):
             idx += 1
-        if (
-            is_multiplexed_enabled(TransactionType.PARTITIONED)
-            and idx < len(requests)
-            and isinstance(requests[idx], CreateSessionRequest)
-        ):
+        # Multiplexed sessions are always enabled - skip CreateSessionRequest
+        if idx < len(requests) and isinstance(requests[idx], CreateSessionRequest):
             idx += 1
         begin_request: BeginTransactionRequest = requests[idx]
         self.assertEqual(

tests/system/test_dbapi.py

@@ -32,9 +32,7 @@
 from google.cloud.spanner_v1 import gapic_version as package_version
 from google.api_core.datetime_helpers import DatetimeWithNanoseconds
 
-from google.cloud.spanner_v1.database_sessions_manager import TransactionType
 from . import _helpers
-from tests._helpers import is_multiplexed_enabled
 
 DATABASE_NAME = "dbapi-txn"
 SPANNER_RPC_PREFIX = "/google.spanner.v1.Spanner/"
@@ -161,20 +159,14 @@ def test_commit(self, client_side):
 
         assert got_rows == [updated_row]
 
-    @pytest.mark.skipif(
-        _helpers.USE_EMULATOR,
-        reason="Emulator does not support multiple parallel transactions.",
+    @pytest.mark.skip(
+        reason="Multiplexed sessions can't be deleted and this test relies on session deletion."
     )
     def test_commit_exception(self):
         """Test that if exception during commit method is caught, then
         subsequent operations on same Cursor and Connection object works
         properly."""
 
-        if is_multiplexed_enabled(transaction_type=TransactionType.READ_WRITE):
-            pytest.skip(
-                "Mutiplexed session can't be deleted and this test relies on session deletion."
-            )
-
         self._execute_common_statements(self._cursor)
         # deleting the session to fail the commit
         self._conn._session.delete()

tests/system/test_observability_options.py

@@ -16,14 +16,12 @@
 from mock import PropertyMock, patch
 
 from google.cloud.spanner_v1.session import Session
-from google.cloud.spanner_v1.database_sessions_manager import TransactionType
 from . import _helpers
 from google.cloud.spanner_v1 import Client
 from google.api_core.exceptions import Aborted
 from google.auth.credentials import AnonymousCredentials
 from google.rpc import code_pb2
 
-from .._helpers import is_multiplexed_enabled
 
 HAS_OTEL_INSTALLED = False
 
@@ -115,18 +113,9 @@ def test_propagation(enable_extended_tracing):
         )  # "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
-        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
-
-        # Determine expected session span name based on multiplexed sessions
-        expected_session_span_name = (
-            "CloudSpanner.CreateMultiplexedSession"
-            if multiplexed_enabled
-            else "CloudSpanner.CreateSession"
-        )
-
+        # Multiplexed sessions are always enabled
         wantNames = [
-            expected_session_span_name,
+            "CloudSpanner.CreateMultiplexedSession",
             "CloudSpanner.Snapshot.execute_sql",
         ]
         assert gotNames == wantNames
@@ -216,7 +205,6 @@ def test_transaction_abort_then_retry_spans(mock_session_id):
     from opentelemetry.trace.status import StatusCode
 
     mock_session_id.return_value = session_id = "session-id"
-    multiplexed = is_multiplexed_enabled(TransactionType.READ_WRITE)
 
     db, trace_exporter = create_db_trace_exporter()
 
@@ -238,60 +226,30 @@ def select_in_txn(txn):
 
     got_statuses, got_events = finished_spans_statuses(trace_exporter)
 
-    # Check for the series of events
-    if multiplexed:
-        # With multiplexed sessions, there are no pool-related events
-        want_events = [
-            ("Creating Session", {}),
-            ("Using session", {"id": session_id, "multiplexed": multiplexed}),
-            ("Returning session", {"id": session_id, "multiplexed": multiplexed}),
-            (
-                "Transaction was aborted in user operation, retrying",
-                {"delay_seconds": "EPHEMERAL", "cause": "EPHEMERAL", "attempt": 1},
-            ),
-            ("Starting Commit", {}),
-            ("Commit Done", {}),
-        ]
-    else:
-        # With regular sessions, include pool-related 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", {}),
-            ("Using session", {"id": session_id, "multiplexed": multiplexed}),
-            ("Returning session", {"id": session_id, "multiplexed": multiplexed}),
-            (
-                "Transaction was aborted in user operation, retrying",
-                {"delay_seconds": "EPHEMERAL", "cause": "EPHEMERAL", "attempt": 1},
-            ),
-            ("Starting Commit", {}),
-            ("Commit Done", {}),
-        ]
+    # Multiplexed sessions are always enabled, no pool-related events
+    want_events = [
+        ("Creating Session", {}),
+        ("Using session", {"id": session_id, "multiplexed": True}),
+        ("Returning session", {"id": session_id, "multiplexed": True}),
+        (
+            "Transaction was aborted in user operation, retrying",
+            {"delay_seconds": "EPHEMERAL", "cause": "EPHEMERAL", "attempt": 1},
+        ),
+        ("Starting Commit", {}),
+        ("Commit Done", {}),
+    ]
     assert got_events == want_events
 
-    # Check for the statues.
+    # Check for the statuses
     codes = StatusCode
-    if multiplexed:
-        # With multiplexed sessions, the session span name is different
-        want_statuses = [
-            ("CloudSpanner.Database.run_in_transaction", codes.OK, None),
-            ("CloudSpanner.CreateMultiplexedSession", codes.OK, None),
-            ("CloudSpanner.Session.run_in_transaction", codes.OK, None),
-            ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
-            ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
-            ("CloudSpanner.Transaction.commit", codes.OK, None),
-        ]
-    else:
-        # With regular sessions
-        want_statuses = [
-            ("CloudSpanner.Database.run_in_transaction", codes.OK, None),
-            ("CloudSpanner.CreateSession", codes.OK, None),
-            ("CloudSpanner.Session.run_in_transaction", codes.OK, None),
-            ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
-            ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
-            ("CloudSpanner.Transaction.commit", codes.OK, None),
-        ]
+    want_statuses = [
+        ("CloudSpanner.Database.run_in_transaction", codes.OK, None),
+        ("CloudSpanner.CreateMultiplexedSession", codes.OK, None),
+        ("CloudSpanner.Session.run_in_transaction", codes.OK, None),
+        ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
+        ("CloudSpanner.Transaction.execute_sql", codes.OK, None),
+        ("CloudSpanner.Transaction.commit", codes.OK, None),
+    ]
     assert got_statuses == want_statuses
 
 
@@ -417,19 +375,10 @@ def tx_update(txn):
     span_list = sorted(span_list, key=lambda span: span.start_time)
     got_span_names = [span.name for span in span_list]
 
-    # Check if multiplexed sessions are enabled for read-write transactions
-    multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
-
-    # Determine expected session span name based on multiplexed sessions
-    expected_session_span_name = (
-        "CloudSpanner.CreateMultiplexedSession"
-        if multiplexed_enabled
-        else "CloudSpanner.CreateSession"
-    )
-
+    # Multiplexed sessions are always enabled
     want_span_names = [
         "CloudSpanner.Database.run_in_transaction",
-        expected_session_span_name,
+        "CloudSpanner.CreateMultiplexedSession",
         "CloudSpanner.Session.run_in_transaction",
         "CloudSpanner.Transaction.commit",
         "CloudSpanner.Transaction.begin",
@@ -462,81 +411,38 @@ def test_database_partitioned_error():
         pass
 
     got_statuses, got_events = finished_spans_statuses(trace_exporter)
-    multiplexed_enabled = is_multiplexed_enabled(TransactionType.PARTITIONED)
-
-    if multiplexed_enabled:
-        expected_event_names = [
-            "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[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
+    # Multiplexed sessions are always enabled
+    expected_event_names = [
+        "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[6][1]["multiplexed"] is False
+    assert got_events[1][1]["multiplexed"] is True
+    assert got_events[3][1]["multiplexed"] is True
 
-        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"]
-            )
+    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"]
 
     codes = StatusCode
-
-    expected_session_span_name = (
-        "CloudSpanner.CreateMultiplexedSession"
-        if multiplexed_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       ^",
         ),
-        (expected_session_span_name, codes.OK, None),
+        ("CloudSpanner.CreateMultiplexedSession", codes.OK, None),
         (
             "CloudSpanner.ExecuteStreamingSql",
             codes.ERROR,

tests/system/test_session_api.py

@@ -421,10 +421,8 @@ def handle_abort(self, database):
         self.handler_done.set()
 
 
+@pytest.mark.skip(reason="Multiplexed sessions do not support CRUD operations.")
 def test_session_crud(sessions_database):
-    if is_multiplexed_enabled(transaction_type=TransactionType.READ_ONLY):
-        pytest.skip("Multiplexed sessions do not support CRUD operations.")
-
     session = sessions_database.session()
     assert not session.exists()
 
@@ -457,33 +455,21 @@ def test_batch_insert_then_read(sessions_database, ot_exporter):
         nth_req0 = sampling_req_id[-2]
 
         db = sessions_database
-        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
         # [A] Verify batch checkout spans
         # -------------------------------
 
         request_id_1 = f"1.{REQ_RAND_PROCESS_ID}.{db._nth_client_id}.{db._channel_id}.{nth_req0 + 0}.1"
 
-        if multiplexed_enabled:
-            assert_span_attributes(
-                ot_exporter,
-                "CloudSpanner.CreateMultiplexedSession",
-                attributes=_make_attributes(
-                    db_name, x_goog_spanner_request_id=request_id_1
-                ),
-                span=span_list[0],
-            )
-        else:
-            assert_span_attributes(
-                ot_exporter,
-                "CloudSpanner.GetSession",
-                attributes=_make_attributes(
-                    db_name,
-                    session_found=True,
-                    x_goog_spanner_request_id=request_id_1,
-                ),
-                span=span_list[0],
-            )
+        # Multiplexed sessions are always enabled
+        assert_span_attributes(
+            ot_exporter,
+            "CloudSpanner.CreateMultiplexedSession",
+            attributes=_make_attributes(
+                db_name, x_goog_spanner_request_id=request_id_1
+            ),
+            span=span_list[0],
+        )
 
         assert_span_attributes(
             ot_exporter,
@@ -500,24 +486,14 @@ def test_batch_insert_then_read(sessions_database, ot_exporter):
         # ----------------------------------
 
         if len(span_list) == 4:
-            if multiplexed_enabled:
-                expected_snapshot_span_name = "CloudSpanner.CreateMultiplexedSession"
-                snapshot_session_attributes = _make_attributes(
-                    db_name,
-                    x_goog_spanner_request_id=f"1.{REQ_RAND_PROCESS_ID}.{db._nth_client_id}.{db._channel_id}.{nth_req0 + 2}.1",
-                )
-            else:
-                expected_snapshot_span_name = "CloudSpanner.GetSession"
-                snapshot_session_attributes = _make_attributes(
-                    db_name,
-                    session_found=True,
-                    x_goog_spanner_request_id=f"1.{REQ_RAND_PROCESS_ID}.{db._nth_client_id}.{db._channel_id}.{nth_req0 + 2}.1",
-                )
-
+            # Multiplexed sessions are always enabled
             assert_span_attributes(
                 ot_exporter,
-                expected_snapshot_span_name,
-                attributes=snapshot_session_attributes,
+                "CloudSpanner.CreateMultiplexedSession",
+                attributes=_make_attributes(
+                    db_name,
+                    x_goog_spanner_request_id=f"1.{REQ_RAND_PROCESS_ID}.{db._nth_client_id}.{db._channel_id}.{nth_req0 + 2}.1",
+                ),
                 span=span_list[2],
             )