diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py index 5e64e106db99..a7d6433530b1 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py @@ -213,6 +213,12 @@ class _SnapshotBase(_SessionWrapper): Allows reuse of API request methods with different transaction selector. + .. note:: + Single-use snapshots (``multi_use=False``) are designed for a single read + or query operation and are thread-confined; they are not safe for concurrent + invocation across multiple threads. Multi-use snapshots and transactions + synchronize concurrent operations using internal locks. + :type session: :class:`~google.cloud.spanner_v1.session.Session` :param session: the session used to perform transaction operations. """ @@ -220,24 +226,30 @@ class _SnapshotBase(_SessionWrapper): _read_only: bool = True _multi_use: bool = False - def __init__(self, session, client_context=None): + def __init__(self, session, client_context=None, multi_use: Optional[bool] = None): super().__init__(session) self._client_context = _validate_client_context(client_context) self._execute_sql_request_count: int = 0 self._read_request_count: int = 0 self._begin_request_sent: bool = False + if multi_use is not None: + self._multi_use = multi_use + # Identifier for the transaction. self._transaction_id: Optional[bytes] = None self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None # Operations within a transaction can be performed concurrently, so we # need to use a lock when updating the transaction. - self._lock: CrossSync.Lock = CrossSync.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: CrossSync.Event = CrossSync.Event() + # Single-use snapshots never execute concurrently or begin transactions, + # so we skip allocating Lock and Event for them. + if self._multi_use: + self._lock: Optional[CrossSync.Lock] = CrossSync.Lock() + self._transaction_begin_event: Optional[CrossSync.Event] = CrossSync.Event() + else: + self._lock = None + self._transaction_begin_event = None @property def _resource_info(self): @@ -258,15 +270,23 @@ async def _wait_for_transaction_begin(self) -> None: id is available, must wait for that first request to complete instead of assuming that the transaction has not begun. + For single-use snapshots, this method checks and enforces sequential + reuse prevention without synchronization, as single-use snapshots are + thread-confined. + :raises ValueError: if the transaction has already been used to execute a request, but is not a multi-use transaction, or if the concurrent request that began the transaction did not complete in time. """ + if not self._multi_use: + if self._begin_request_sent or self._read_request_count > 0: + raise ValueError("Cannot re-use single-use snapshot.") + self._begin_request_sent = True + return + async with self._lock: if self._begin_request_sent or self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") wait_needed = self._transaction_id is None else: wait_needed = False @@ -663,7 +683,7 @@ async def _get_streamed_result_set( trace_name = f"CloudSpanner.{type(self).__name__}.{trace_method_name}" is_inline_begin = False - if self._transaction_id is None: + if self._multi_use and self._transaction_id is None: is_inline_begin = True await self._lock.acquire() @@ -983,7 +1003,8 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: if self._transaction_id is None and transaction_pb.id: self._transaction_id = transaction_pb.id # Release any request waiting for this transaction to begin. - self._transaction_begin_event.set() + if self._transaction_begin_event is not None: + self._transaction_begin_event.set() if transaction_pb._pb.HasField("precommit_token"): self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token) @@ -993,7 +1014,10 @@ async def _update_for_precommit_token_pb( self, precommit_token_pb: MultiplexedSessionPrecommitToken ) -> None: """Updates the snapshot for the given multiplexed session precommit token.""" - async with self._lock: + if self._lock is not None: + async with self._lock: + self._update_for_precommit_token_pb_unsafe(precommit_token_pb) + else: self._update_for_precommit_token_pb_unsafe(precommit_token_pb) def _update_for_precommit_token_pb_unsafe( @@ -1029,7 +1053,9 @@ def __init__( transaction_id=None, client_context=None, ): - super(Snapshot, self).__init__(session, client_context=client_context) + super(Snapshot, self).__init__( + session, client_context=client_context, multi_use=multi_use + ) opts = [read_timestamp, min_read_timestamp, max_staleness, exact_staleness] flagged = [opt for opt in opts if opt is not None] if len(flagged) > 1: @@ -1047,7 +1073,6 @@ def __init__( self._min_read_timestamp = min_read_timestamp self._max_staleness = max_staleness self._exact_staleness = exact_staleness - self._multi_use = multi_use self._transaction_id = transaction_id def _build_transaction_options_pb(self) -> TransactionOptions: diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py index 425857e45095..f72988d5c8c7 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py @@ -84,7 +84,9 @@ class Transaction(_SnapshotBase, _BatchBase): _read_only: bool = False def __init__(self, session, client_context=None): - super(Transaction, self).__init__(session, client_context=client_context) + super(Transaction, self).__init__( + session, client_context=client_context, multi_use=True + ) self.rolled_back: bool = False # If this transaction is used to retry a previous aborted transaction with a diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py index ce2bd07306d8..32967dab775e 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -187,6 +187,12 @@ class _SnapshotBase(_SessionWrapper): Allows reuse of API request methods with different transaction selector. + .. note:: + Single-use snapshots (``multi_use=False``) are designed for a single read + or query operation and are thread-confined; they are not safe for concurrent + invocation across multiple threads. Multi-use snapshots and transactions + synchronize concurrent operations using internal locks. + :type session: :class:`~google.cloud.spanner_v1.session.Session` :param session: the session used to perform transaction operations. """ @@ -194,18 +200,26 @@ class _SnapshotBase(_SessionWrapper): _read_only: bool = True _multi_use: bool = False - def __init__(self, session, client_context=None): + def __init__(self, session, client_context=None, multi_use: Optional[bool] = None): super().__init__(session) self._client_context = _validate_client_context(client_context) self._execute_sql_request_count: int = 0 self._read_request_count: int = 0 self._begin_request_sent: bool = False + if multi_use is not None: + self._multi_use = multi_use self._transaction_id: Optional[bytes] = None self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None - self._lock: CrossSync._Sync_Impl.Lock = CrossSync._Sync_Impl.Lock() - self._transaction_begin_event: CrossSync._Sync_Impl.Event = ( - CrossSync._Sync_Impl.Event() - ) + if self._multi_use: + self._lock: Optional[CrossSync._Sync_Impl.Lock] = ( + CrossSync._Sync_Impl.Lock() + ) + self._transaction_begin_event: Optional[CrossSync._Sync_Impl.Event] = ( + CrossSync._Sync_Impl.Event() + ) + else: + self._lock = None + self._transaction_begin_event = None @property def _resource_info(self): @@ -225,13 +239,20 @@ def _wait_for_transaction_begin(self) -> None: id is available, must wait for that first request to complete instead of assuming that the transaction has not begun. + For single-use snapshots, this method checks and enforces sequential + reuse prevention without synchronization, as single-use snapshots are + thread-confined. + :raises ValueError: if the transaction has already been used to execute a request, but is not a multi-use transaction, or if the concurrent request that began the transaction did not complete in time.""" + if not self._multi_use: + if self._begin_request_sent or self._read_request_count > 0: + raise ValueError("Cannot re-use single-use snapshot.") + self._begin_request_sent = True + return with self._lock: if self._begin_request_sent or self._read_request_count > 0: - if not self._multi_use: - raise ValueError("Cannot re-use single-use snapshot.") wait_needed = self._transaction_id is None else: wait_needed = False @@ -596,7 +617,7 @@ def _get_streamed_result_set( trace_method_name = "execute_sql" if is_execute_sql_request else "read" trace_name = f"CloudSpanner.{type(self).__name__}.{trace_method_name}" is_inline_begin = False - if self._transaction_id is None: + if self._multi_use and self._transaction_id is None: is_inline_begin = True self._lock.acquire() try: @@ -882,7 +903,8 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None: """Updates the snapshot for the given transaction.""" if self._transaction_id is None and transaction_pb.id: self._transaction_id = transaction_pb.id - self._transaction_begin_event.set() + if self._transaction_begin_event is not None: + self._transaction_begin_event.set() if transaction_pb._pb.HasField("precommit_token"): self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token) @@ -890,7 +912,10 @@ def _update_for_precommit_token_pb( self, precommit_token_pb: MultiplexedSessionPrecommitToken ) -> None: """Updates the snapshot for the given multiplexed session precommit token.""" - with self._lock: + if self._lock is not None: + with self._lock: + self._update_for_precommit_token_pb_unsafe(precommit_token_pb) + else: self._update_for_precommit_token_pb_unsafe(precommit_token_pb) def _update_for_precommit_token_pb_unsafe( @@ -918,7 +943,9 @@ def __init__( transaction_id=None, client_context=None, ): - super(Snapshot, self).__init__(session, client_context=client_context) + super(Snapshot, self).__init__( + session, client_context=client_context, multi_use=multi_use + ) opts = [read_timestamp, min_read_timestamp, max_staleness, exact_staleness] flagged = [opt for opt in opts if opt is not None] if len(flagged) > 1: @@ -934,7 +961,6 @@ def __init__( self._min_read_timestamp = min_read_timestamp self._max_staleness = max_staleness self._exact_staleness = exact_staleness - self._multi_use = multi_use self._transaction_id = transaction_id def _build_transaction_options_pb(self) -> TransactionOptions: diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py index 1bd8d01d348d..4d0dd1c9d2c3 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py @@ -73,7 +73,9 @@ class Transaction(_SnapshotBase, _BatchBase): _read_only: bool = False def __init__(self, session, client_context=None): - super(Transaction, self).__init__(session, client_context=client_context) + super(Transaction, self).__init__( + session, client_context=client_context, multi_use=True + ) self.rolled_back: bool = False self._multiplexed_session_previous_transaction_id: Optional[bytes] = None diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py index e4e8785f2c43..329f1b748b08 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -24,15 +24,21 @@ ServiceUnavailable, ) -from google.cloud.spanner_v1._async.snapshot import Snapshot +from google.cloud.spanner_v1._async.snapshot import Snapshot, _SnapshotBase +from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken from google.cloud.spanner_v1.types.result_set import PartialResultSet, ResultSetMetadata from google.cloud.spanner_v1.types.spanner import ( ExecuteSqlRequest, Partition, PartitionResponse, ) -from google.cloud.spanner_v1.types.transaction import Transaction as TransactionPB -from google.cloud.spanner_v1.types.transaction import TransactionSelector +from google.cloud.spanner_v1.types.transaction import ( + Transaction as TransactionPB, +) +from google.cloud.spanner_v1.types.transaction import ( + TransactionOptions, + TransactionSelector, +) from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode TABLE_NAME = "citizens" @@ -43,6 +49,16 @@ DURATION = timedelta(seconds=3) +class _Derived(_SnapshotBase): + """A minimally-implemented _SnapshotBase-derived class for testing.""" + + transaction_tag = None + TRANSACTION_OPTIONS = TransactionOptions() + + def _build_transaction_options_pb(self) -> TransactionOptions: + return self.TRANSACTION_OPTIONS + + class Test_snapshot_coverage(unittest.IsolatedAsyncioTestCase): def setUp(self): self.patch_metrics = mock.patch( @@ -61,13 +77,81 @@ def tearDown(self): self.patch_trace.stop() def _make_snapshot(self, *args, **kwargs): - s = Snapshot(*args, **kwargs) - # FORCE _lock to exist if it doesn't (though constructor should handle it) - if not hasattr(s, "_lock"): - from google.cloud.aio._cross_sync.cross_sync import CrossSync + return Snapshot(*args, **kwargs) + + def test_derived_constructor(self): + session = _Session() + derived = _Derived(session=session) + self.assertTrue(derived._read_only) + self.assertFalse(derived._multi_use) + self.assertIsNone(derived._lock) + self.assertIsNone(derived._transaction_begin_event) + + def test_derived_constructor_multi_use(self): + session = _Session() + derived = _Derived(session=session, multi_use=True) + self.assertTrue(derived._read_only) + self.assertTrue(derived._multi_use) + self.assertIsInstance(derived._lock, type(asyncio.Lock())) + self.assertIsInstance(derived._transaction_begin_event, type(asyncio.Event())) + + def test_ctor_single_use_no_lock(self): + snapshot = self._make_snapshot(_Session(), multi_use=False) + self.assertFalse(snapshot._multi_use) + self.assertIsNone(snapshot._lock) + self.assertIsNone(snapshot._transaction_begin_event) + + def test_ctor_multi_use(self): + snapshot = self._make_snapshot(_Session(), multi_use=True) + self.assertTrue(snapshot._multi_use) + self.assertIsInstance(snapshot._lock, type(asyncio.Lock())) + self.assertIsInstance(snapshot._transaction_begin_event, type(asyncio.Event())) + + async def test_wait_for_transaction_begin_twice_single_use_raises(self): + snapshot = self._make_snapshot(_Session(), multi_use=False) + await snapshot._wait_for_transaction_begin() + self.assertTrue(snapshot._begin_request_sent) + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + await snapshot._wait_for_transaction_begin() + + async def test_update_for_precommit_token_pb_multi_use(self): + token = MultiplexedSessionPrecommitToken(seq_num=1) + snapshot = self._make_snapshot(_Session(), multi_use=True) + await snapshot._update_for_precommit_token_pb(token) + self.assertEqual(snapshot._precommit_token, token) - s._lock = CrossSync.Lock() - return s + async def test_update_for_precommit_token_pb_single_use(self): + token = MultiplexedSessionPrecommitToken(seq_num=1) + snapshot = self._make_snapshot(_Session(), multi_use=False) + await snapshot._update_for_precommit_token_pb(token) + self.assertEqual(snapshot._precommit_token, token) + + async def test_execute_sql_twice_single_use_fails(self): + session = _Session() + session._database.spanner_api.execute_streaming_sql.return_value = ( + _MockIterator(PartialResultSet()) + ) + snapshot = self._make_snapshot(session, multi_use=False) + result_set = await snapshot.execute_sql("SELECT 1") + async for _ in result_set: + pass + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + await snapshot.execute_sql("SELECT 1") + + async def test_read_twice_single_use_fails(self): + from google.cloud.spanner_v1.keyset import KeySet + + session = _Session() + session._database.spanner_api.streaming_read.return_value = _MockIterator( + PartialResultSet() + ) + snapshot = self._make_snapshot(session, multi_use=False) + keyset = KeySet(all_=True) + result_set = await snapshot.read(TABLE_NAME, COLUMNS, keyset) + async for _ in result_set: + pass + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + await snapshot.read(TABLE_NAME, COLUMNS, keyset) async def test_read_errors(self): snapshot = self._make_snapshot(_Session(), multi_use=False) @@ -140,11 +224,20 @@ async def test_partition_read_ok(self): self.assertEqual(tokens, [token_1]) def test__update_for_transaction_pb(self): - snapshot = self._make_snapshot(_Session()) + snapshot = self._make_snapshot(_Session(), multi_use=False) + pb = TransactionPB(id=TXN_ID, read_timestamp=TIMESTAMP) + snapshot._update_for_transaction_pb(pb) + self.assertEqual(snapshot._transaction_id, TXN_ID) + self.assertEqual(snapshot._transaction_read_timestamp, TIMESTAMP) + + def test__update_for_transaction_pb_multi_use(self): + snapshot = self._make_snapshot(_Session(), multi_use=True) + self.assertFalse(snapshot._transaction_begin_event.is_set()) pb = TransactionPB(id=TXN_ID, read_timestamp=TIMESTAMP) snapshot._update_for_transaction_pb(pb) self.assertEqual(snapshot._transaction_id, TXN_ID) self.assertEqual(snapshot._transaction_read_timestamp, TIMESTAMP) + self.assertTrue(snapshot._transaction_begin_event.is_set()) async def test_restart_on_unavailable_precommit(self): from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index 4a65d46aea8c..0ef57c7b1c9b 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -476,8 +476,7 @@ def test_iteration_w_raw_w_multiuse(self): database = _Database() database.spanner_api = build_spanner_api() session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = True + derived = _build_snapshot_derived(session, multi_use=True) resumable = self._call_fut(derived, restart, request, session=session) self.assertEqual(list(resumable), list(FIRST)) self.assertEqual(len(restart.mock_calls), 1) @@ -506,8 +505,7 @@ def test_iteration_w_raw_raising_unavailable_w_multiuse(self): database = _Database() database.spanner_api = build_spanner_api() session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = True + derived = _build_snapshot_derived(session, multi_use=True) resumable = self._call_fut(derived, restart, request, session=session) self.assertEqual(list(resumable), list(SECOND)) self.assertEqual(len(restart.mock_calls), 2) @@ -541,8 +539,7 @@ def test_iteration_w_raw_raising_unavailable_after_token_w_multiuse(self): database = _Database() database.spanner_api = build_spanner_api() session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = True + derived = _build_snapshot_derived(session, multi_use=True) resumable = self._call_fut(derived, restart, request, session=session) @@ -685,7 +682,7 @@ def test_iteration_w_multiple_span_creation(self, mock_region): class Test_SnapshotBase(OpenTelemetryBase): def test_ctor(self): session = build_session() - derived = _build_snapshot_derived(session=session) + derived = _Derived(session=session) # Attributes from _SessionWrapper. self.assertIs(derived._session, session) @@ -698,11 +695,20 @@ def test_ctor(self): self.assertFalse(derived._begin_request_sent) self.assertIsNone(derived._transaction_id) self.assertIsNone(derived._precommit_token) - self.assertIsInstance(derived._lock, type(Lock())) - self.assertFalse(derived._transaction_begin_event.is_set()) + self.assertIsNone(derived._lock) + self.assertIsNone(derived._transaction_begin_event) self.assertNoSpans() + def test_derived_constructor_multi_use(self): + session = build_session() + derived = _Derived(session=session, multi_use=True) + + self.assertTrue(derived._read_only) + self.assertTrue(derived._multi_use) + self.assertIsInstance(derived._lock, type(Lock())) + self.assertFalse(derived._transaction_begin_event.is_set()) + def test__wait_for_transaction_begin_claims_inline_begin(self): derived = _build_snapshot_derived(multi_use=True) @@ -718,6 +724,14 @@ def test__wait_for_transaction_begin_wo_multi_use(self): with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): derived._wait_for_transaction_begin() + def test__wait_for_transaction_begin_twice_wo_multi_use(self): + derived = _build_snapshot_derived(multi_use=False) + derived._wait_for_transaction_begin() + self.assertTrue(derived._begin_request_sent) + + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + derived._wait_for_transaction_begin() + def test__wait_for_transaction_begin_waits_for_concurrent_begin(self): """A concurrent request must wait for the in-flight inline begin. @@ -1067,8 +1081,7 @@ def _execute_read( api = database.spanner_api = build_spanner_api() api.streaming_read.return_value = _MockIterator(*result_sets) session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = multi_use + derived = _build_snapshot_derived(session, multi_use=multi_use) derived._read_request_count = count if not first: @@ -1733,8 +1746,7 @@ def _partition_read_helper( api = database.spanner_api = build_spanner_api() api.partition_read.return_value = response session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = multi_use + derived = _build_snapshot_derived(session, multi_use=multi_use) if w_txn: derived._transaction_id = TXN_ID @@ -1859,8 +1871,7 @@ def test_partition_read_w_retry(self): ] session = _Session(database) - derived = _build_snapshot_derived(session) - derived._multi_use = True + derived = _build_snapshot_derived(session, multi_use=True) derived._transaction_id = TXN_ID list(derived.partition_read(TABLE_NAME, COLUMNS, keyset)) @@ -2130,7 +2141,8 @@ def test_ctor_defaults(self): self.assertEqual(snapshot._read_request_count, 0) self.assertIsNone(snapshot._transaction_id) self.assertIsNone(snapshot._precommit_token) - self.assertIsInstance(snapshot._lock, type(Lock())) + self.assertIsNone(snapshot._lock) + self.assertIsNone(snapshot._transaction_begin_event) # Attributes from Snapshot. self.assertTrue(snapshot._strong) @@ -2139,6 +2151,96 @@ def test_ctor_defaults(self): self.assertIsNone(snapshot._max_staleness) self.assertIsNone(snapshot._exact_staleness) + def test_ctor_multi_use(self): + session = build_session() + snapshot = build_snapshot(session=session, multi_use=True) + + self.assertIs(snapshot._session, session) + self.assertTrue(snapshot._read_only) + self.assertTrue(snapshot._multi_use) + self.assertIsNone(snapshot._transaction_id) + self.assertIsInstance(snapshot._lock, type(Lock())) + self.assertIsInstance(snapshot._transaction_begin_event, type(Event())) + + def test_ctor_single_use_no_lock(self): + session = build_session() + snapshot = build_snapshot(session=session, multi_use=False) + self.assertFalse(snapshot._multi_use) + self.assertIsNone(snapshot._lock) + self.assertIsNone(snapshot._transaction_begin_event) + + def test_update_for_precommit_token_pb_multi_use(self): + token = build_precommit_token_pb(seq_num=1) + snapshot = self._make_one(_Session(), multi_use=True) + snapshot._update_for_precommit_token_pb(token) + self.assertEqual(snapshot._precommit_token, token) + + def test_update_for_precommit_token_pb_single_use(self): + token = build_precommit_token_pb(seq_num=1) + snapshot = self._make_one(_Session(), multi_use=False) + snapshot._update_for_precommit_token_pb(token) + self.assertEqual(snapshot._precommit_token, token) + + def test__update_for_transaction_pb(self): + from datetime import timezone + + from google.cloud.spanner_v1.types.transaction import ( + Transaction as TransactionPB, + ) + + timestamp = datetime.now(timezone.utc) + snapshot = self._make_one(_Session(), multi_use=False) + pb = TransactionPB(id=TXN_ID, read_timestamp=timestamp) + snapshot._update_for_transaction_pb(pb) + self.assertEqual(snapshot._transaction_id, TXN_ID) + self.assertEqual(snapshot._transaction_read_timestamp, timestamp) + + def test__update_for_transaction_pb_multi_use(self): + from datetime import timezone + + from google.cloud.spanner_v1.types.transaction import ( + Transaction as TransactionPB, + ) + + timestamp = datetime.now(timezone.utc) + snapshot = self._make_one(_Session(), multi_use=True) + self.assertFalse(snapshot._transaction_begin_event.is_set()) + pb = TransactionPB(id=TXN_ID, read_timestamp=timestamp) + snapshot._update_for_transaction_pb(pb) + self.assertEqual(snapshot._transaction_id, TXN_ID) + self.assertEqual(snapshot._transaction_read_timestamp, timestamp) + self.assertTrue(snapshot._transaction_begin_event.is_set()) + + def test_execute_sql_twice_single_use_fails(self): + from google.cloud.spanner_v1 import PartialResultSet + + database = _Database() + database.spanner_api = build_spanner_api() + database.spanner_api.execute_streaming_sql.return_value = _MockIterator( + PartialResultSet() + ) + session = _Session(database) + snapshot = self._make_one(session=session, multi_use=False) + list(snapshot.execute_sql(SQL_QUERY)) + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + list(snapshot.execute_sql(SQL_QUERY)) + + def test_read_twice_single_use_fails(self): + from google.cloud.spanner_v1 import PartialResultSet + from google.cloud.spanner_v1.keyset import KeySet + + database = _Database() + database.spanner_api = build_spanner_api() + database.spanner_api.streaming_read.return_value = _MockIterator( + PartialResultSet() + ) + session = _Session(database) + snapshot = self._make_one(session=session, multi_use=False) + keyset = KeySet(all_=True) + list(snapshot.read(TABLE_NAME, COLUMNS, keyset)) + with self.assertRaisesRegex(ValueError, "Cannot re-use single-use snapshot."): + list(snapshot.read(TABLE_NAME, COLUMNS, keyset)) + def test_ctor_w_multiple_options(self): with self.assertRaises(ValueError): build_snapshot(read_timestamp=datetime.min, max_staleness=timedelta()) @@ -2373,8 +2475,7 @@ def _build_snapshot_derived(session=None, multi_use=False, read_only=True) -> _D if session.session_id is None: session._session_id = "session-id" - derived = _Derived(session=session) - derived._multi_use = multi_use + derived = _Derived(session=session, multi_use=multi_use) derived._read_only = read_only return derived