fix lint

Author: rahul2393

google/cloud/spanner_v1/database.py

@@ -868,9 +868,7 @@ def session(self, labels=None, database_role=None):
         # instead.
         role = database_role or self._database_role
         # Always use multiplexed sessions
-        return Session(
-            self, labels=labels, database_role=role, is_multiplexed=True
-        )
+        return Session(self, labels=labels, database_role=role, is_multiplexed=True)
 
     def snapshot(self, **kw):
         """Return an object which wraps a snapshot.

google/cloud/spanner_v1/database_sessions_manager.py

@@ -210,4 +210,3 @@ def _maintain_multiplexed_session(session_manager_ref) -> None:
                 manager._multiplexed_session = manager._build_multiplexed_session()
 
             session_created_time = time()
-

tests/_helpers.py

@@ -1,5 +1,4 @@
 import unittest
-from os import getenv
 
 import mock
 
@@ -36,21 +35,12 @@
 
 
 def is_multiplexed_enabled(transaction_type: TransactionType) -> bool:
-    """Returns whether multiplexed sessions are enabled for the given transaction type."""
+    """Returns whether multiplexed sessions are enabled for the given transaction type.
 
-    env_var = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
-    env_var_partitioned = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
-    env_var_read_write = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
-
-    def _getenv(val: str) -> bool:
-        return getenv(val, "true").lower().strip() != "false"
-
-    if transaction_type is TransactionType.READ_ONLY:
-        return _getenv(env_var)
-    elif transaction_type is TransactionType.PARTITIONED:
-        return _getenv(env_var) and _getenv(env_var_partitioned)
-    else:
-        return _getenv(env_var) and _getenv(env_var_read_write)
+    Multiplexed sessions are now always enabled for all transaction types.
+    This function is kept for backward compatibility with existing tests.
+    """
+    return True
 
 
 def get_test_ot_exporter():

tests/mockserver_tests/mock_server_test_base.py

@@ -41,7 +41,6 @@
     SpannerServicer,
     start_mock_server,
 )
-from tests._helpers import is_multiplexed_enabled
 
 
 # Creates an aborted status with the smallest possible retry delay.
@@ -240,52 +239,30 @@ def assert_requests_sequence(
         transaction_type,
         allow_multiple_batch_create=True,
     ):
-        """Assert that the requests sequence matches the expected types, accounting for multiplexed sessions and retries.
+        """Assert that the requests sequence matches the expected types, accounting for multiplexed sessions.
 
         Args:
             requests: List of requests from spanner_service.requests
             expected_types: List of expected request types (excluding session creation requests)
-            transaction_type: TransactionType enum value to check multiplexed session status
-            allow_multiple_batch_create: If True, skip all leading BatchCreateSessionsRequest and one optional CreateSessionRequest
+            transaction_type: TransactionType enum value (unused, kept for backward compatibility)
+            allow_multiple_batch_create: If True, skip leading CreateSessionRequest (kept for backward compatibility)
         """
-        from google.cloud.spanner_v1 import (
-            BatchCreateSessionsRequest,
-            CreateSessionRequest,
-        )
+        from google.cloud.spanner_v1 import CreateSessionRequest
 
-        mux_enabled = is_multiplexed_enabled(transaction_type)
         idx = 0
-        # Skip all leading BatchCreateSessionsRequest (for retries)
+        # Skip CreateSessionRequest for multiplexed session
         if allow_multiple_batch_create:
             while idx < len(requests) and isinstance(
-                requests[idx], BatchCreateSessionsRequest
-            ):
-                idx += 1
-            # For multiplexed, optionally skip a CreateSessionRequest
-            if (
-                mux_enabled
-                and idx < len(requests)
-                and isinstance(requests[idx], CreateSessionRequest)
+                requests[idx], CreateSessionRequest
             ):
                 idx += 1
         else:
-            if mux_enabled:
-                self.assertTrue(
-                    isinstance(requests[idx], BatchCreateSessionsRequest),
-                    f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}",
-                )
-                idx += 1
-                self.assertTrue(
-                    isinstance(requests[idx], CreateSessionRequest),
-                    f"Expected CreateSessionRequest at index {idx}, got {type(requests[idx])}",
-                )
-                idx += 1
-            else:
-                self.assertTrue(
-                    isinstance(requests[idx], BatchCreateSessionsRequest),
-                    f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}",
-                )
-                idx += 1
+            # Expect exactly one CreateSessionRequest for multiplexed session
+            self.assertTrue(
+                isinstance(requests[idx], CreateSessionRequest),
+                f"Expected CreateSessionRequest at index {idx}, got {type(requests[idx])}",
+            )
+            idx += 1
         # Check the rest of the expected request types
         for expected_type in expected_types:
             self.assertTrue(
@@ -303,13 +280,12 @@ def adjust_request_id_sequence(self, expected_segments, requests, transaction_ty
         Args:
             expected_segments: List of expected (method, (sequence_numbers)) tuples
             requests: List of actual requests from spanner_service.requests
-            transaction_type: TransactionType enum value to check multiplexed session status
+            transaction_type: TransactionType enum value (unused, kept for backward compatibility)
 
         Returns:
             List of adjusted expected segments with corrected sequence numbers
         """
         from google.cloud.spanner_v1 import (
-            BatchCreateSessionsRequest,
             CreateSessionRequest,
             ExecuteSqlRequest,
             BeginTransactionRequest,
@@ -318,15 +294,13 @@ def adjust_request_id_sequence(self, expected_segments, requests, transaction_ty
         # Count session creation requests that come before the first non-session request
         session_requests_before = 0
         for req in requests:
-            if isinstance(req, (BatchCreateSessionsRequest, CreateSessionRequest)):
+            if isinstance(req, CreateSessionRequest):
                 session_requests_before += 1
             elif isinstance(req, (ExecuteSqlRequest, BeginTransactionRequest)):
                 break
 
-        # For multiplexed sessions, we expect 2 session requests (BatchCreateSessions + CreateSession)
-        # For non-multiplexed, we expect 1 session request (BatchCreateSessions)
-        mux_enabled = is_multiplexed_enabled(transaction_type)
-        expected_session_requests = 2 if mux_enabled else 1
+        # With multiplexed sessions, we expect 1 session request (CreateSession)
+        expected_session_requests = 1
         extra_session_requests = session_requests_before - expected_session_requests
 
         # Adjust sequence numbers based on extra session requests

tests/mockserver_tests/test_request_id_header.py

@@ -16,7 +16,6 @@
 import threading
 
 from google.cloud.spanner_v1 import (
-    BatchCreateSessionsRequest,
     CreateSessionRequest,
     ExecuteSqlRequest,
     BeginTransactionRequest,
@@ -58,20 +57,17 @@ def test_snapshot_execute_sql(self):
         NTH_CLIENT = self.database._nth_client_id
         CHANNEL_ID = self.database._channel_id
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
-        # Filter out CreateSessionRequest unary segments for comparison
-        filtered_unary_segments = [
-            seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")
-        ]
+        # With multiplexed sessions, we expect one CreateSession request
         want_unary_segments = [
             (
-                "/google.spanner.v1.Spanner/BatchCreateSessions",
+                "/google.spanner.v1.Spanner/CreateSession",
                 (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 1, 1),
             )
         ]
         # Dynamically determine the expected sequence number for ExecuteStreamingSql
         session_requests_before = 0
         for req in requests:
-            if isinstance(req, (BatchCreateSessionsRequest, CreateSessionRequest)):
+            if isinstance(req, CreateSessionRequest):
                 session_requests_before += 1
             elif isinstance(req, ExecuteSqlRequest):
                 break
@@ -88,7 +84,7 @@ def test_snapshot_execute_sql(self):
                 ),
             )
         ]
-        assert filtered_unary_segments == want_unary_segments
+        assert got_unary_segments == want_unary_segments
         assert got_stream_segments == want_stream_segments
 
     def test_snapshot_read_concurrent(self):
@@ -118,45 +114,32 @@ def select1():
         for thread in threads:
             thread.join()
         requests = self.spanner_service.requests
-        # Allow for an extra request due to multiplexed session creation
-        expected_min = 2 + n
-        expected_max = expected_min + 1
+        # With multiplexed sessions: 1 CreateSession + (n + 1) ExecuteSql
+        expected_min = 1 + n + 1
+        expected_max = expected_min
         assert (
             expected_min <= len(requests) <= expected_max
-        ), f"Expected {expected_min} or {expected_max} requests, got {len(requests)}: {requests}"
+        ), f"Expected {expected_min} requests, got {len(requests)}: {requests}"
         client_id = db._nth_client_id
         channel_id = db._channel_id
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
         want_unary_segments = [
             (
-                "/google.spanner.v1.Spanner/BatchCreateSessions",
+                "/google.spanner.v1.Spanner/CreateSession",
                 (1, REQ_RAND_PROCESS_ID, client_id, channel_id, 1, 1),
             ),
         ]
         assert any(seg == want_unary_segments[0] for seg in got_unary_segments)
 
-        # Dynamically determine the expected sequence numbers for ExecuteStreamingSql
-        session_requests_before = 0
-        for req in requests:
-            if isinstance(req, (BatchCreateSessionsRequest, CreateSessionRequest)):
-                session_requests_before += 1
-            elif isinstance(req, ExecuteSqlRequest):
-                break
-        want_stream_segments = [
-            (
-                "/google.spanner.v1.Spanner/ExecuteStreamingSql",
-                (
-                    1,
-                    REQ_RAND_PROCESS_ID,
-                    client_id,
-                    channel_id,
-                    session_requests_before + i,
-                    1,
-                ),
-            )
-            for i in range(1, n + 2)
-        ]
-        assert got_stream_segments == want_stream_segments
+        # Verify we have the expected number of ExecuteStreamingSql segments
+        # (n + 1 = 11 for initial + 10 concurrent)
+        assert len(got_stream_segments) == n + 1
+        # Verify all segments are for ExecuteStreamingSql
+        for seg in got_stream_segments:
+            assert seg[0] == "/google.spanner.v1.Spanner/ExecuteStreamingSql"
+            # Verify the segment has correct client_id and channel_id
+            assert seg[1][2] == client_id
+            assert seg[1][3] == channel_id
 
     def test_database_run_in_transaction_retries_on_abort(self):
         counters = dict(aborted=0)
@@ -192,33 +175,22 @@ def test_database_execute_partitioned_dml_request_id(self):
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
         NTH_CLIENT = self.database._nth_client_id
         CHANNEL_ID = self.database._channel_id
-        # Allow for extra unary segments due to session creation
-        filtered_unary_segments = [
-            seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")
-        ]
         # Find the actual sequence number for BeginTransaction
         begin_txn_seq = None
-        for seg in filtered_unary_segments:
+        for seg in got_unary_segments:
             if seg[0].endswith("/BeginTransaction"):
                 begin_txn_seq = seg[1][4]
                 break
         want_unary_segments = [
             (
-                "/google.spanner.v1.Spanner/BatchCreateSessions",
+                "/google.spanner.v1.Spanner/CreateSession",
                 (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 1, 1),
             ),
             (
                 "/google.spanner.v1.Spanner/BeginTransaction",
                 (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, begin_txn_seq, 1),
             ),
         ]
-        # Dynamically determine the expected sequence number for ExecuteStreamingSql
-        session_requests_before = 0
-        for req in requests:
-            if isinstance(req, (BatchCreateSessionsRequest, CreateSessionRequest)):
-                session_requests_before += 1
-            elif isinstance(req, ExecuteSqlRequest):
-                break
         # Find the actual sequence number for ExecuteStreamingSql
         exec_sql_seq = got_stream_segments[0][1][4] if got_stream_segments else None
         want_stream_segments = [
@@ -227,12 +199,12 @@ def test_database_execute_partitioned_dml_request_id(self):
                 (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, exec_sql_seq, 1),
             )
         ]
-        assert all(seg in filtered_unary_segments for seg in want_unary_segments)
+        assert all(seg in got_unary_segments for seg in want_unary_segments)
         assert got_stream_segments == want_stream_segments
 
     def test_unary_retryable_error(self):
         add_select1_result()
-        add_error(SpannerServicer.BatchCreateSessions.__name__, unavailable_status())
+        add_error(SpannerServicer.CreateSession.__name__, unavailable_status())
 
         if not getattr(self.database, "_interceptors", None):
             self.database._interceptors = MockServerTestBase._interceptors

tests/mockserver_tests/test_tags.py

@@ -23,7 +23,6 @@
     MockServerTestBase,
     add_single_result,
 )
-from tests._helpers import is_multiplexed_enabled
 from google.cloud.spanner_v1.database_sessions_manager import TransactionType
 
 
@@ -100,8 +99,9 @@ def test_select_read_only_transaction_with_transaction_tag(self):
             TransactionType.READ_ONLY,
         )
         # Transaction tags are not supported for read-only transactions.
-        mux_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
-        tag_idx = 3 if mux_enabled else 2
+        # With multiplexed sessions: CreateSession, BeginTransaction, ExecuteSql, ExecuteSql
+        # ExecuteSql requests start at index 2
+        tag_idx = 2
         self.assertEqual("", requests[tag_idx].request_options.transaction_tag)
         self.assertEqual("", requests[tag_idx + 1].request_options.transaction_tag)
 
@@ -155,8 +155,9 @@ def test_select_read_write_transaction_with_transaction_tag(self):
             ],
             TransactionType.READ_WRITE,
         )
-        mux_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
-        tag_idx = 3 if mux_enabled else 2
+        # With multiplexed sessions: CreateSession, BeginTransaction, ExecuteSql, ExecuteSql, Commit
+        # ExecuteSql requests start at index 2, Commit at index 4
+        tag_idx = 2
         self.assertEqual(
             "my_transaction_tag", requests[tag_idx].request_options.transaction_tag
         )
@@ -187,8 +188,9 @@ def test_select_read_write_transaction_with_transaction_and_request_tag(self):
             ],
             TransactionType.READ_WRITE,
         )
-        mux_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
-        tag_idx = 3 if mux_enabled else 2
+        # With multiplexed sessions: CreateSession, BeginTransaction, ExecuteSql, ExecuteSql, Commit
+        # ExecuteSql requests start at index 2, Commit at index 4
+        tag_idx = 2
         self.assertEqual(
             "my_transaction_tag", requests[tag_idx].request_options.transaction_tag
         )

tests/system/test_database_api.py

@@ -249,9 +249,7 @@ def test_update_ddl_w_operation_id(
     #    https://github.com/GoogleCloudPlatform/google-cloud-python/issues/5629
     # )
     temp_db_id = _helpers.unique_id("update_ddl", separator="_")
-    temp_db = shared_instance.database(
-        temp_db_id, database_dialect=database_dialect
-    )
+    temp_db = shared_instance.database(temp_db_id, database_dialect=database_dialect)
     create_op = temp_db.create()
     databases_to_delete.append(temp_db)
     create_op.result(DBAPI_OPERATION_TIMEOUT)  # raises on failure / timeout.