chore: updates thread handling to avoid race condition in tests

Author: chalmerlowe

packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py

@@ -220,6 +220,7 @@ def __init__(self, session, client_context=None):
         # respectively.
         self._execute_sql_request_count: int = 0
         self._read_request_count: int = 0
+        self._begin_request_sent: bool = False
 
         # Identifier for the transaction.
         self._transaction_id: Optional[bytes] = None
@@ -229,10 +230,14 @@ def __init__(self, session, client_context=None):
         # highest sequence number is included in the  commit request.
         self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None
 
-        # Operations within a transaction can be performed using multiple
+        # Operation within a transaction can be performed using multiple
         # threads, so we need to use a lock when updating the transaction.
         self._lock: threading.Lock = threading.Lock()
 
+        # Event to coordinate concurrent requests beginning the transaction.
+        # This is used to prevent the "Transaction has not begun" race condition.
+        self._transaction_begin_event = threading.Event()
+
     def begin(self) -> bytes:
         """Begins a transaction on the database.
 
@@ -341,11 +346,27 @@ def read(
             read request, but is not a multi-use transaction or has not begun.
         """
 
-        if self._read_request_count > 0:
-            if not self._multi_use:
-                raise ValueError("Cannot re-use single-use snapshot.")
-            if self._transaction_id is None:
-                raise ValueError("Transaction has not begun.")
+        with self._lock:
+            # Check if this request is beginning the transaction.
+            # If a request is already in progress, other requests must wait
+            # until the transaction ID is available.
+            if self._begin_request_sent or self._read_request_count > 0:
+                if not self._multi_use:
+                    raise ValueError("Cannot re-use single-use snapshot.")
+                if self._transaction_id is None:
+                    wait_needed = True
+                else:
+                    wait_needed = False
+            else:
+                wait_needed = False
+                self._begin_request_sent = True
+
+        if wait_needed:
+            # Wait for the transaction to begin (set by another concurrent request).
+            # This prevents the race condition where concurrent requests think
+            # the transaction hasn't begun.
+            if not self._transaction_begin_event.wait(timeout=30.0):
+                raise ValueError("Timed out waiting for transaction to begin.")
 
         session = self._session
         database = session._database
@@ -527,11 +548,27 @@ def execute_sql(
             read request, but is not a multi-use transaction or has not begun.
         """
 
-        if self._read_request_count > 0:
-            if not self._multi_use:
-                raise ValueError("Cannot re-use single-use snapshot.")
-            if self._transaction_id is None:
-                raise ValueError("Transaction has not begun.")
+        with self._lock:
+            # Check if this request is beginning the transaction.
+            # If a request is already in progress, other requests must wait
+            # until the transaction ID is available.
+            if self._begin_request_sent or self._read_request_count > 0:
+                if not self._multi_use:
+                    raise ValueError("Cannot re-use single-use snapshot.")
+                if self._transaction_id is None:
+                    wait_needed = True
+                else:
+                    wait_needed = False
+            else:
+                wait_needed = False
+                self._begin_request_sent = True
+
+        if wait_needed:
+            # Wait for the transaction to begin (set by another concurrent request).
+            # This prevents the race condition where concurrent requests think
+            # the transaction hasn't begun.
+            if not self._transaction_begin_event.wait(timeout=30.0):
+                raise ValueError("Timed out waiting for transaction to begin.")
 
         if params is not None:
             params_pb = Struct(
@@ -1058,6 +1095,8 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None:
         # caller is responsible for locking until the transaction ID is updated.
         if self._transaction_id is None and transaction_pb.id:
             self._transaction_id = transaction_pb.id
+            # Notify waiting threads that the transaction has begun.
+            self._transaction_begin_event.set()
 
         if transaction_pb._pb.HasField("precommit_token"):
             self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token)

packages/google-cloud-spanner/tests/unit/test_spanner.py

@@ -1100,6 +1100,7 @@ def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_
     def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_read(
         self,
     ):
+        self.maxDiff = None
         database = _Database()
         api = database.spanner_api = self._make_spanner_api()
         session = _Session(database)