fix systests

Author: rahul2393

google/cloud/spanner_v1/database.py

@@ -848,7 +848,14 @@ def session(self, labels=None, database_role=None):
         # If role is specified in param, then that role is used
         # instead.
         role = database_role or self._database_role
-        return Session(self, labels=labels, database_role=role)
+        is_multiplexed = False
+        if self.sessions_manager._use_multiplexed(
+            transaction_type=TransactionType.READ_ONLY
+        ):
+            is_multiplexed = True
+        return Session(
+            self, labels=labels, database_role=role, is_multiplexed=is_multiplexed
+        )
 
     def snapshot(self, **kw):
         """Return an object which wraps a snapshot.

google/cloud/spanner_v1/session.py

@@ -275,7 +275,13 @@ def delete(self):
                 current_span, "Deleting Session failed due to unset session_id"
             )
             raise ValueError("Session ID not set by back-end")
-
+        if self._is_multiplexed:
+            add_span_event(
+                current_span,
+                "Skipped deleting Multiplexed Session",
+                {"session.id": self._session_id},
+            )
+            return
         add_span_event(
             current_span, "Deleting Session", {"session.id": self._session_id}
         )

google/cloud/spanner_v1/snapshot.py

@@ -133,6 +133,8 @@ def _restart_on_unavailable(
                 # Update the transaction from the response.
                 if transaction is not None:
                     transaction._update_for_result_set_pb(item)
+                if item.precommit_token is not None and transaction is not None:
+                    transaction._update_for_precommit_token_pb(item.precommit_token)
 
                 if item.resume_token:
                     resume_token = item.resume_token
@@ -1013,9 +1015,6 @@ def _update_for_result_set_pb(
         if result_set_pb.metadata and result_set_pb.metadata.transaction:
             self._update_for_transaction_pb(result_set_pb.metadata.transaction)
 
-        if result_set_pb.precommit_token:
-            self._update_for_precommit_token_pb(result_set_pb.precommit_token)
-
     def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None:
         """Updates the snapshot for the given transaction.
 
@@ -1031,7 +1030,7 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None:
             self._transaction_id = transaction_pb.id
 
         if transaction_pb.precommit_token:
-            self._update_for_precommit_token_pb(transaction_pb.precommit_token)
+            self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token)
 
     def _update_for_precommit_token_pb(
         self, precommit_token_pb: MultiplexedSessionPrecommitToken
@@ -1044,10 +1043,22 @@ def _update_for_precommit_token_pb(
         # Because multiple threads can be used to perform operations within a
         # transaction, we need to use a lock when updating the precommit token.
         with self._lock:
-            if self._precommit_token is None or (
-                precommit_token_pb.seq_num > self._precommit_token.seq_num
-            ):
-                self._precommit_token = precommit_token_pb
+            self._update_for_precommit_token_pb_unsafe(precommit_token_pb)
+
+    def _update_for_precommit_token_pb_unsafe(
+        self, precommit_token_pb: MultiplexedSessionPrecommitToken
+    ) -> None:
+        """Updates the snapshot for the given multiplexed session precommit token.
+        This method is unsafe because it does not acquire a lock before updating
+        the precommit token. It should only be used when the caller has already
+        acquired the lock.
+        :type precommit_token_pb: :class:`~google.cloud.spanner_v1.MultiplexedSessionPrecommitToken`
+        :param precommit_token_pb: The multiplexed session precommit token to update the snapshot with.
+        """
+        if self._precommit_token is None or (
+            precommit_token_pb.seq_num > self._precommit_token.seq_num
+        ):
+            self._precommit_token = precommit_token_pb
 
 
 class Snapshot(_SnapshotBase):

google/cloud/spanner_v1/transaction.py

@@ -516,6 +516,9 @@ def wrapped_method(*args, **kwargs):
         if is_inline_begin:
             self._lock.release()
 
+        if result_set_pb.precommit_token is not None:
+            self._update_for_precommit_token_pb(result_set_pb.precommit_token)
+
         return result_set_pb.stats.row_count_exact
 
     def batch_update(
@@ -660,6 +663,9 @@ def wrapped_method(*args, **kwargs):
         if is_inline_begin:
             self._lock.release()
 
+        if len(response_pb.result_sets) > 0 and response_pb.result_sets[0].precommit_token:
+            self._update_for_precommit_token_pb(response_pb.result_sets[0].precommit_token)
+
         row_counts = [
             result_set.stats.row_count_exact for result_set in response_pb.result_sets
         ]
@@ -736,9 +742,6 @@ def _update_for_execute_batch_dml_response_pb(
         :type response_pb: :class:`~google.cloud.spanner_v1.types.ExecuteBatchDmlResponse`
         :param response_pb: The execute batch DML response to update the transaction with.
         """
-        if response_pb.precommit_token:
-            self._update_for_precommit_token_pb(response_pb.precommit_token)
-
         # Only the first result set contains the result set metadata.
         if len(response_pb.result_sets) > 0:
             self._update_for_result_set_pb(response_pb.result_sets[0])

tests/system/_sample_data.py

@@ -101,7 +101,6 @@ def _assert_timestamp(value, nano_value):
             # Allow up to 1 microsecond difference for timestamp precision issues
             us_diff = abs(value.microsecond - nano_value.microsecond)
             if us_diff > 1:
-                print(f"DEBUG: Microsecond comparison failed:")
                 print(f"  Expected: {value} (microsecond: {value.microsecond})")
                 print(f"  Found: {nano_value} (microsecond: {nano_value.microsecond})")
                 print(f"  Difference: {us_diff} microseconds")

tests/system/test_session_api.py

@@ -424,6 +424,9 @@ def handle_abort(self, database):
 
 
 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()
 
@@ -1412,11 +1415,13 @@ def unit_of_work(transaction):
     for span in ot_exporter.get_finished_spans():
         if span and span.name:
             span_list.append(span)
-
+    multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
     span_list = sorted(span_list, key=lambda v1: v1.start_time)
     got_span_names = [span.name for span in span_list]
     expected_span_names = [
-        "CloudSpanner.CreateSession",
+        "CloudSpanner.CreateMultiplexedSession"
+        if multiplexed_enabled
+        else "CloudSpanner.CreateSession",
         "CloudSpanner.Batch.commit",
         "Test Span",
         "CloudSpanner.Session.run_in_transaction",