From cdaf25b35d27355e4ea577843004fdc2d16bb4ac Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 29 Sep 2023 12:47:51 -0700 Subject: [PATCH] fix: ensure transactions rollback on failure (#767) --- .../cloud/firestore_v1/async_transaction.py | 85 ++-- google/cloud/firestore_v1/base_transaction.py | 3 - google/cloud/firestore_v1/transaction.py | 86 ++--- google/cloud/firestore_v1/watch.py | 4 +- tests/unit/v1/test_async_transaction.py | 357 ++++++++--------- tests/unit/v1/test_transaction.py | 364 +++++++++--------- 6 files changed, 406 insertions(+), 493 deletions(-) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index f4ecf32d3..b504bebad 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -110,6 +110,7 @@ async def _rollback(self) -> None: Raises: ValueError: If no transaction is in progress. + google.api_core.exceptions.GoogleAPICallError: If the rollback fails. """ if not self.in_progress: raise ValueError(_CANT_ROLLBACK) @@ -124,6 +125,7 @@ async def _rollback(self) -> None: metadata=self._client._rpc_metadata, ) finally: + # clean up, even if rollback fails self._clean_up() async def _commit(self) -> list: @@ -223,10 +225,6 @@ async def _pre_commit( ) -> Coroutine: """Begin transaction and call the wrapped coroutine. - If the coroutine raises an exception, the transaction will be rolled - back. If not, the transaction will be "ready" for ``Commit`` (i.e. - it will have staged writes). - Args: transaction (:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`): @@ -250,41 +248,7 @@ async def _pre_commit( self.current_id = transaction._id if self.retry_id is None: self.retry_id = self.current_id - try: - return await self.to_wrap(transaction, *args, **kwargs) - except: # noqa - # NOTE: If ``rollback`` fails this will lose the information - # from the original failure. - await transaction._rollback() - raise - - async def _maybe_commit(self, transaction: AsyncTransaction) -> bool: - """Try to commit the transaction. - - If the transaction is read-write and the ``Commit`` fails with the - ``ABORTED`` status code, it will be retried. Any other failure will - not be caught. - - Args: - transaction - (:class:`~google.cloud.firestore_v1.transaction.Transaction`): - The transaction to be ``Commit``-ed. - - Returns: - bool: Indicating if the commit succeeded. - """ - try: - await transaction._commit() - return True - except exceptions.GoogleAPICallError as exc: - if transaction._read_only: - raise - - if isinstance(exc, exceptions.Aborted): - # If a read-write transaction returns ABORTED, retry. - return False - else: - raise + return await self.to_wrap(transaction, *args, **kwargs) async def __call__(self, transaction, *args, **kwargs): """Execute the wrapped callable within a transaction. @@ -306,22 +270,35 @@ async def __call__(self, transaction, *args, **kwargs): ``max_attempts``. """ self._reset() + retryable_exceptions = ( + (exceptions.Aborted) if not transaction._read_only else () + ) + last_exc = None - for attempt in range(transaction._max_attempts): - result = await self._pre_commit(transaction, *args, **kwargs) - succeeded = await self._maybe_commit(transaction) - if succeeded: - return result - - # Subsequent requests will use the failed transaction ID as part of - # the ``BeginTransactionRequest`` when restarting this transaction - # (via ``options.retry_transaction``). This preserves the "spot in - # line" of the transaction, so exponential backoff is not required - # in this case. - - await transaction._rollback() - msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) - raise ValueError(msg) + try: + for attempt in range(transaction._max_attempts): + result = await self._pre_commit(transaction, *args, **kwargs) + try: + await transaction._commit() + return result + except retryable_exceptions as exc: + last_exc = exc + # Retry attempts that result in retryable exceptions + # Subsequent requests will use the failed transaction ID as part of + # the ``BeginTransactionRequest`` when restarting this transaction + # (via ``options.retry_transaction``). This preserves the "spot in + # line" of the transaction, so exponential backoff is not required + # in this case. + # retries exhausted + # wrap the last exception in a ValueError before raising + msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) + raise ValueError(msg) from last_exc + + except BaseException: + # rollback the transaction on any error + # errors raised during _rollback will be chained to the original error through __context__ + await transaction._rollback() + raise def async_transactional( diff --git a/google/cloud/firestore_v1/base_transaction.py b/google/cloud/firestore_v1/base_transaction.py index 145321245..b4e5dd038 100644 --- a/google/cloud/firestore_v1/base_transaction.py +++ b/google/cloud/firestore_v1/base_transaction.py @@ -185,8 +185,5 @@ def _reset(self) -> None: def _pre_commit(self, transaction, *args, **kwargs) -> NoReturn: raise NotImplementedError - def _maybe_commit(self, transaction) -> NoReturn: - raise NotImplementedError - def __call__(self, transaction, *args, **kwargs): raise NotImplementedError diff --git a/google/cloud/firestore_v1/transaction.py b/google/cloud/firestore_v1/transaction.py index cfcb968c8..3c175a4ce 100644 --- a/google/cloud/firestore_v1/transaction.py +++ b/google/cloud/firestore_v1/transaction.py @@ -44,7 +44,7 @@ # Types needed only for Type Hints from google.cloud.firestore_v1.base_document import DocumentSnapshot from google.cloud.firestore_v1.types import CommitResponse -from typing import Any, Callable, Generator, Optional +from typing import Any, Callable, Generator class Transaction(batch.WriteBatch, BaseTransaction): @@ -108,6 +108,7 @@ def _rollback(self) -> None: Raises: ValueError: If no transaction is in progress. + google.api_core.exceptions.GoogleAPICallError: If the rollback fails. """ if not self.in_progress: raise ValueError(_CANT_ROLLBACK) @@ -122,6 +123,7 @@ def _rollback(self) -> None: metadata=self._client._rpc_metadata, ) finally: + # clean up, even if rollback fails self._clean_up() def _commit(self) -> list: @@ -214,10 +216,6 @@ def __init__(self, to_wrap) -> None: def _pre_commit(self, transaction: Transaction, *args, **kwargs) -> Any: """Begin transaction and call the wrapped callable. - If the callable raises an exception, the transaction will be rolled - back. If not, the transaction will be "ready" for ``Commit`` (i.e. - it will have staged writes). - Args: transaction (:class:`~google.cloud.firestore_v1.transaction.Transaction`): @@ -241,41 +239,7 @@ def _pre_commit(self, transaction: Transaction, *args, **kwargs) -> Any: self.current_id = transaction._id if self.retry_id is None: self.retry_id = self.current_id - try: - return self.to_wrap(transaction, *args, **kwargs) - except: # noqa - # NOTE: If ``rollback`` fails this will lose the information - # from the original failure. - transaction._rollback() - raise - - def _maybe_commit(self, transaction: Transaction) -> Optional[bool]: - """Try to commit the transaction. - - If the transaction is read-write and the ``Commit`` fails with the - ``ABORTED`` status code, it will be retried. Any other failure will - not be caught. - - Args: - transaction - (:class:`~google.cloud.firestore_v1.transaction.Transaction`): - The transaction to be ``Commit``-ed. - - Returns: - bool: Indicating if the commit succeeded. - """ - try: - transaction._commit() - return True - except exceptions.GoogleAPICallError as exc: - if transaction._read_only: - raise - - if isinstance(exc, exceptions.Aborted): - # If a read-write transaction returns ABORTED, retry. - return False - else: - raise + return self.to_wrap(transaction, *args, **kwargs) def __call__(self, transaction: Transaction, *args, **kwargs): """Execute the wrapped callable within a transaction. @@ -297,22 +261,34 @@ def __call__(self, transaction: Transaction, *args, **kwargs): ``max_attempts``. """ self._reset() + retryable_exceptions = ( + (exceptions.Aborted) if not transaction._read_only else () + ) + last_exc = None - for attempt in range(transaction._max_attempts): - result = self._pre_commit(transaction, *args, **kwargs) - succeeded = self._maybe_commit(transaction) - if succeeded: - return result - - # Subsequent requests will use the failed transaction ID as part of - # the ``BeginTransactionRequest`` when restarting this transaction - # (via ``options.retry_transaction``). This preserves the "spot in - # line" of the transaction, so exponential backoff is not required - # in this case. - - transaction._rollback() - msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) - raise ValueError(msg) + try: + for attempt in range(transaction._max_attempts): + result = self._pre_commit(transaction, *args, **kwargs) + try: + transaction._commit() + return result + except retryable_exceptions as exc: + last_exc = exc + # Retry attempts that result in retryable exceptions + # Subsequent requests will use the failed transaction ID as part of + # the ``BeginTransactionRequest`` when restarting this transaction + # (via ``options.retry_transaction``). This preserves the "spot in + # line" of the transaction, so exponential backoff is not required + # in this case. + # retries exhausted + # wrap the last exception in a ValueError before raising + msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) + raise ValueError(msg) from last_exc + except BaseException: # noqa: B901 + # rollback the transaction on any error + # errors raised during _rollback will be chained to the original error through __context__ + transaction._rollback() + raise def transactional(to_wrap: Callable) -> _Transactional: diff --git a/google/cloud/firestore_v1/watch.py b/google/cloud/firestore_v1/watch.py index d1ce5a57a..eabc218de 100644 --- a/google/cloud/firestore_v1/watch.py +++ b/google/cloud/firestore_v1/watch.py @@ -401,7 +401,9 @@ def _on_snapshot_target_change_remove(self, target_change): error_message = "Error %s: %s" % (code, message) - raise RuntimeError(error_message) + raise RuntimeError(error_message) from exceptions.from_grpc_status( + code, message + ) def _on_snapshot_target_change_reset(self, target_change): # Whatever changes have happened so far no longer matter. diff --git a/tests/unit/v1/test_async_transaction.py b/tests/unit/v1/test_async_transaction.py index 12f704a6e..7c1ab0650 100644 --- a/tests/unit/v1/test_async_transaction.py +++ b/tests/unit/v1/test_async_transaction.py @@ -158,7 +158,6 @@ async def test_asynctransaction__rollback_not_allowed(): with pytest.raises(ValueError) as exc_info: await transaction._rollback() - assert exc_info.value.args == (_CANT_ROLLBACK,) @@ -460,135 +459,147 @@ async def test_asynctransactional__pre_commit_retry_id_already_set_success(): @pytest.mark.asyncio -async def test_asynctransactional__pre_commit_failure(): - exc = RuntimeError("Nope not today.") - to_wrap = AsyncMock(side_effect=exc, spec=[]) +async def test_asynctransactional___call__success_first_attempt(): + to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) wrapped = _make_async_transactional(to_wrap) - txn_id = b"gotta-fail" + txn_id = b"whole-enchilada" transaction = _make_transaction(txn_id) - with pytest.raises(RuntimeError) as exc_info: - await wrapped._pre_commit(transaction, 10, 20) - assert exc_info.value is exc + result = await wrapped(transaction, "a", b="c") + assert result is mock.sentinel.result assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - to_wrap.assert_called_once_with(transaction, 10, 20) + to_wrap.assert_called_once_with(transaction, "a", b="c") firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, metadata=transaction._client._rpc_metadata, ) - firestore_api.rollback.assert_called_once_with( + firestore_api.rollback.assert_not_called() + firestore_api.commit.assert_called_once_with( request={ "database": transaction._client._database_string, + "writes": [], "transaction": txn_id, }, metadata=transaction._client._rpc_metadata, ) - firestore_api.commit.assert_not_called() @pytest.mark.asyncio -async def test_asynctransactional__pre_commit_failure_with_rollback_failure(): +async def test_asynctransactional___call__success_second_attempt(): from google.api_core import exceptions + from google.cloud.firestore_v1.types import common + from google.cloud.firestore_v1.types import firestore + from google.cloud.firestore_v1.types import write - exc1 = ValueError("I will not be only failure.") - to_wrap = AsyncMock(side_effect=exc1, spec=[]) + to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) wrapped = _make_async_transactional(to_wrap) - txn_id = b"both-will-fail" + txn_id = b"whole-enchilada" transaction = _make_transaction(txn_id) - # Actually force the ``rollback`` to fail as well. - exc2 = exceptions.InternalServerError("Rollback blues.") + + # Actually force the ``commit`` to fail on first / succeed on second. + exc = exceptions.Aborted("Contention junction.") firestore_api = transaction._client._firestore_api - firestore_api.rollback.side_effect = exc2 + firestore_api.commit.side_effect = [ + exc, + firestore.CommitResponse(write_results=[write.WriteResult()]), + ] - # Try to ``_pre_commit`` - with pytest.raises(exceptions.InternalServerError) as exc_info: - await wrapped._pre_commit(transaction, a="b", c="zebra") - assert exc_info.value is exc2 + # Call the __call__-able ``wrapped``. + result = await wrapped(transaction, "a", b="c") + assert result is mock.sentinel.result assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - to_wrap.assert_called_once_with(transaction, a="b", c="zebra") - firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, - ) - firestore_api.rollback.assert_called_once_with( - request={ - "database": transaction._client._database_string, - "transaction": txn_id, - }, - metadata=transaction._client._rpc_metadata, - ) - firestore_api.commit.assert_not_called() - - -@pytest.mark.asyncio -async def test_asynctransactional__maybe_commit_success(): - wrapped = _make_async_transactional(mock.sentinel.callable_) - - txn_id = b"nyet" - transaction = _make_transaction(txn_id) - transaction._id = txn_id # We won't call ``begin()``. - succeeded = await wrapped._maybe_commit(transaction) - assert succeeded - - # On success, _id is reset. - assert transaction._id is None - - # Verify mocks. + wrapped_call = mock.call(transaction, "a", b="c") + assert to_wrap.mock_calls == [wrapped_call, wrapped_call] firestore_api = transaction._client._firestore_api - firestore_api.begin_transaction.assert_not_called() + db_str = transaction._client._database_string + options_ = common.TransactionOptions( + read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + ) + expected_calls = [ + mock.call( + request={"database": db_str, "options": None}, + metadata=transaction._client._rpc_metadata, + ), + mock.call( + request={"database": db_str, "options": options_}, + metadata=transaction._client._rpc_metadata, + ), + ] + assert firestore_api.begin_transaction.mock_calls == expected_calls firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( - request={ - "database": transaction._client._database_string, - "writes": [], - "transaction": txn_id, - }, + commit_call = mock.call( + request={"database": db_str, "writes": [], "transaction": txn_id}, metadata=transaction._client._rpc_metadata, ) + assert firestore_api.commit.mock_calls == [commit_call, commit_call] +@pytest.mark.parametrize("max_attempts", [1, 5]) @pytest.mark.asyncio -async def test_asynctransactional__maybe_commit_failure_read_only(): +async def test_asynctransactional___call__failure_max_attempts(max_attempts): + """ + rasie retryable error and exhause max_attempts + """ from google.api_core import exceptions + from google.cloud.firestore_v1.types import common + from google.cloud.firestore_v1.async_transaction import _EXCEED_ATTEMPTS_TEMPLATE - wrapped = _make_async_transactional(mock.sentinel.callable_) + to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) + wrapped = _make_async_transactional(to_wrap) - txn_id = b"failed" - transaction = _make_transaction(txn_id, read_only=True) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. + txn_id = b"attempt_exhaustion" + transaction = _make_transaction(txn_id, max_attempts=max_attempts) - # Actually force the ``commit`` to fail (use ABORTED, but cannot - # retry since read-only). - exc = exceptions.Aborted("Read-only did a bad.") + # Actually force the ``commit`` to fail. + exc = exceptions.Aborted("Contention just once.") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc - with pytest.raises(exceptions.Aborted) as exc_info: - await wrapped._maybe_commit(transaction) - assert exc_info.value is exc + # Call the __call__-able ``wrapped``. + with pytest.raises(ValueError) as exc_info: + await wrapped(transaction, "here", there=1.5) - assert transaction._id == txn_id + err_msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) + assert exc_info.value.args == (err_msg,) + # should retain cause exception + assert exc_info.value.__cause__ == exc + + assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + assert to_wrap.call_count == max_attempts + to_wrap.assert_called_with(transaction, "here", there=1.5) + assert firestore_api.begin_transaction.call_count == max_attempts + options_ = common.TransactionOptions( + read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + ) + expected_calls = [ + mock.call( + request={ + "database": transaction._client._database_string, + "options": None if i == 0 else options_, + }, + metadata=transaction._client._rpc_metadata, + ) + for i in range(max_attempts) + ] + assert firestore_api.begin_transaction.call_args_list == expected_calls + assert firestore_api.commit.call_count == max_attempts + firestore_api.commit.assert_called_with( request={ "database": transaction._client._database_string, "writes": [], @@ -596,105 +607,63 @@ async def test_asynctransactional__maybe_commit_failure_read_only(): }, metadata=transaction._client._rpc_metadata, ) - - -@pytest.mark.asyncio -async def test_asynctransactional__maybe_commit_failure_can_retry(): - from google.api_core import exceptions - - wrapped = _make_async_transactional(mock.sentinel.callable_) - - txn_id = b"failed-but-retry" - transaction = _make_transaction(txn_id) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. - - # Actually force the ``commit`` to fail. - exc = exceptions.Aborted("Read-write did a bad.") - firestore_api = transaction._client._firestore_api - firestore_api.commit.side_effect = exc - - succeeded = await wrapped._maybe_commit(transaction) - assert not succeeded - - assert transaction._id == txn_id - assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id - - # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, - "writes": [], "transaction": txn_id, }, metadata=transaction._client._rpc_metadata, ) +@pytest.mark.parametrize("max_attempts", [1, 5]) @pytest.mark.asyncio -async def test_asynctransactional__maybe_commit_failure_cannot_retry(): +async def test_asynctransactional___call__failure_readonly(max_attempts): + """ + readonly transaction should never retry + """ from google.api_core import exceptions + from google.cloud.firestore_v1.types import common - wrapped = _make_async_transactional(mock.sentinel.callable_) + to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) + wrapped = _make_async_transactional(to_wrap) - txn_id = b"failed-but-not-retryable" - transaction = _make_transaction(txn_id) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. + txn_id = b"read_only_fail" + transaction = _make_transaction(txn_id, max_attempts=max_attempts, read_only=True) # Actually force the ``commit`` to fail. - exc = exceptions.InternalServerError("Real bad thing") + exc = exceptions.Aborted("Contention just once.") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc - with pytest.raises(exceptions.InternalServerError) as exc_info: - await wrapped._maybe_commit(transaction) - assert exc_info.value is exc + # Call the __call__-able ``wrapped``. + with pytest.raises(exceptions.Aborted) as exc_info: + await wrapped(transaction, "here", there=1.5) - assert transaction._id == txn_id + assert exc_info.value == exc + + assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + to_wrap.assert_called_once_with(transaction, "here", there=1.5) + firestore_api.begin_transaction.assert_called_once_with( request={ "database": transaction._client._database_string, - "writes": [], - "transaction": txn_id, + "options": common.TransactionOptions( + read_only=common.TransactionOptions.ReadOnly() + ), }, metadata=transaction._client._rpc_metadata, ) - - -@pytest.mark.asyncio -async def test_asynctransactional___call__success_first_attempt(): - to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) - wrapped = _make_async_transactional(to_wrap) - - txn_id = b"whole-enchilada" - transaction = _make_transaction(txn_id) - result = await wrapped(transaction, "a", b="c") - assert result is mock.sentinel.result - - assert transaction._id is None - assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id - - # Verify mocks. - to_wrap.assert_called_once_with(transaction, "a", b="c") - firestore_api = transaction._client._firestore_api - firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, + firestore_api.rollback.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "transaction": txn_id, + }, metadata=transaction._client._rpc_metadata, ) - firestore_api.rollback.assert_not_called() firestore_api.commit.assert_called_once_with( request={ "database": transaction._client._database_string, @@ -705,93 +674,101 @@ async def test_asynctransactional___call__success_first_attempt(): ) +@pytest.mark.parametrize("max_attempts", [1, 5]) @pytest.mark.asyncio -async def test_asynctransactional___call__success_second_attempt(): +async def test_asynctransactional___call__failure_with_non_retryable(max_attempts): + """ + call fails due to an exception that is not retryable. + Should rollback raise immediately + """ from google.api_core import exceptions - from google.cloud.firestore_v1.types import common - from google.cloud.firestore_v1.types import firestore - from google.cloud.firestore_v1.types import write to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) wrapped = _make_async_transactional(to_wrap) - txn_id = b"whole-enchilada" - transaction = _make_transaction(txn_id) + txn_id = b"non_retryable" + transaction = _make_transaction(txn_id, max_attempts=max_attempts) - # Actually force the ``commit`` to fail on first / succeed on second. - exc = exceptions.Aborted("Contention junction.") + # Actually force the ``commit`` to fail. + exc = exceptions.InvalidArgument("non retryable") firestore_api = transaction._client._firestore_api - firestore_api.commit.side_effect = [ - exc, - firestore.CommitResponse(write_results=[write.WriteResult()]), - ] + firestore_api.commit.side_effect = exc # Call the __call__-able ``wrapped``. - result = await wrapped(transaction, "a", b="c") - assert result is mock.sentinel.result + with pytest.raises(exceptions.InvalidArgument) as exc_info: + await wrapped(transaction, "here", there=1.5) + + assert exc_info.value == exc assert transaction._id is None assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id # Verify mocks. - wrapped_call = mock.call(transaction, "a", b="c") - assert to_wrap.mock_calls == [wrapped_call, wrapped_call] - firestore_api = transaction._client._firestore_api - db_str = transaction._client._database_string - options_ = common.TransactionOptions( - read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + to_wrap.assert_called_once_with(transaction, "here", there=1.5) + firestore_api.begin_transaction.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "options": None, + }, + metadata=transaction._client._rpc_metadata, ) - expected_calls = [ - mock.call( - request={"database": db_str, "options": None}, - metadata=transaction._client._rpc_metadata, - ), - mock.call( - request={"database": db_str, "options": options_}, - metadata=transaction._client._rpc_metadata, - ), - ] - assert firestore_api.begin_transaction.mock_calls == expected_calls - firestore_api.rollback.assert_not_called() - commit_call = mock.call( - request={"database": db_str, "writes": [], "transaction": txn_id}, + firestore_api.rollback.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "transaction": txn_id, + }, + metadata=transaction._client._rpc_metadata, + ) + firestore_api.commit.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "writes": [], + "transaction": txn_id, + }, metadata=transaction._client._rpc_metadata, ) - assert firestore_api.commit.mock_calls == [commit_call, commit_call] @pytest.mark.asyncio -async def test_asynctransactional___call__failure(): +async def test_asynctransactional___call__failure_with_rollback_failure(): + """ + Test second failure as part of rollback + should maintain first failure as __context__ + """ from google.api_core import exceptions - from google.cloud.firestore_v1.async_transaction import _EXCEED_ATTEMPTS_TEMPLATE to_wrap = AsyncMock(return_value=mock.sentinel.result, spec=[]) wrapped = _make_async_transactional(to_wrap) - txn_id = b"only-one-shot" + txn_id = b"non_retryable" transaction = _make_transaction(txn_id, max_attempts=1) # Actually force the ``commit`` to fail. - exc = exceptions.Aborted("Contention just once.") + exc = exceptions.InvalidArgument("first error") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc + # also force a second error on rollback + rb_exc = exceptions.InternalServerError("second error") + firestore_api.rollback.side_effect = rb_exc # Call the __call__-able ``wrapped``. - with pytest.raises(ValueError) as exc_info: + # should raise second error with first error as __context__ + with pytest.raises(exceptions.InternalServerError) as exc_info: await wrapped(transaction, "here", there=1.5) - err_msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) - assert exc_info.value.args == (err_msg,) + assert exc_info.value == rb_exc + assert exc_info.value.__context__ == exc assert transaction._id is None assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id # Verify mocks. to_wrap.assert_called_once_with(transaction, "here", there=1.5) firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, + request={ + "database": transaction._client._database_string, + "options": None, + }, metadata=transaction._client._rpc_metadata, ) firestore_api.rollback.assert_called_once_with( diff --git a/tests/unit/v1/test_transaction.py b/tests/unit/v1/test_transaction.py index 27366b276..26bb5cc9c 100644 --- a/tests/unit/v1/test_transaction.py +++ b/tests/unit/v1/test_transaction.py @@ -464,135 +464,149 @@ def test__transactional__pre_commit_retry_id_already_set_success(database): @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__pre_commit_failure(database): - exc = RuntimeError("Nope not today.") - to_wrap = mock.Mock(side_effect=exc, spec=[]) +def test__transactional___call__success_first_attempt(database): + to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) wrapped = _make__transactional(to_wrap) - txn_id = b"gotta-fail" + txn_id = b"whole-enchilada" transaction = _make_transaction_pb(txn_id, database=database) - with pytest.raises(RuntimeError) as exc_info: - wrapped._pre_commit(transaction, 10, 20) - assert exc_info.value is exc + result = wrapped(transaction, "a", b="c") + assert result is mock.sentinel.result assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - to_wrap.assert_called_once_with(transaction, 10, 20) + to_wrap.assert_called_once_with(transaction, "a", b="c") firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, metadata=transaction._client._rpc_metadata, ) - firestore_api.rollback.assert_called_once_with( + firestore_api.rollback.assert_not_called() + firestore_api.commit.assert_called_once_with( request={ "database": transaction._client._database_string, + "writes": [], "transaction": txn_id, }, metadata=transaction._client._rpc_metadata, ) - firestore_api.commit.assert_not_called() @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__pre_commit_failure_with_rollback_failure(database): +def test__transactional___call__success_second_attempt(database): from google.api_core import exceptions + from google.cloud.firestore_v1.types import common + from google.cloud.firestore_v1.types import firestore + from google.cloud.firestore_v1.types import write - exc1 = ValueError("I will not be only failure.") - to_wrap = mock.Mock(side_effect=exc1, spec=[]) + to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) wrapped = _make__transactional(to_wrap) - txn_id = b"both-will-fail" + txn_id = b"whole-enchilada" transaction = _make_transaction_pb(txn_id, database=database) - # Actually force the ``rollback`` to fail as well. - exc2 = exceptions.InternalServerError("Rollback blues.") + + # Actually force the ``commit`` to fail on first / succeed on second. + exc = exceptions.Aborted("Contention junction.") firestore_api = transaction._client._firestore_api - firestore_api.rollback.side_effect = exc2 + firestore_api.commit.side_effect = [ + exc, + firestore.CommitResponse(write_results=[write.WriteResult()]), + ] - # Try to ``_pre_commit`` - with pytest.raises(exceptions.InternalServerError) as exc_info: - wrapped._pre_commit(transaction, a="b", c="zebra") - assert exc_info.value is exc2 + # Call the __call__-able ``wrapped``. + result = wrapped(transaction, "a", b="c") + assert result is mock.sentinel.result assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - to_wrap.assert_called_once_with(transaction, a="b", c="zebra") - firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, - ) - firestore_api.rollback.assert_called_once_with( - request={ - "database": transaction._client._database_string, - "transaction": txn_id, - }, - metadata=transaction._client._rpc_metadata, - ) - firestore_api.commit.assert_not_called() - - -@pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__maybe_commit_success(database): - wrapped = _make__transactional(mock.sentinel.callable_) - - txn_id = b"nyet" - transaction = _make_transaction_pb(txn_id, database=database) - transaction._id = txn_id # We won't call ``begin()``. - succeeded = wrapped._maybe_commit(transaction) - assert succeeded - - # On success, _id is reset. - assert transaction._id is None - - # Verify mocks. + wrapped_call = mock.call(transaction, "a", b="c") + assert to_wrap.mock_calls, [wrapped_call == wrapped_call] firestore_api = transaction._client._firestore_api - firestore_api.begin_transaction.assert_not_called() + db_str = transaction._client._database_string + options_ = common.TransactionOptions( + read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + ) + expected_calls = [ + mock.call( + request={"database": db_str, "options": None}, + metadata=transaction._client._rpc_metadata, + ), + mock.call( + request={"database": db_str, "options": options_}, + metadata=transaction._client._rpc_metadata, + ), + ] + assert firestore_api.begin_transaction.mock_calls == expected_calls firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( - request={ - "database": transaction._client._database_string, - "writes": [], - "transaction": txn_id, - }, + commit_call = mock.call( + request={"database": db_str, "writes": [], "transaction": txn_id}, metadata=transaction._client._rpc_metadata, ) + assert firestore_api.commit.mock_calls == [commit_call, commit_call] @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__maybe_commit_failure_read_only(database): +@pytest.mark.parametrize("max_attempts", [1, 5]) +def test_transactional___call__failure_max_attempts(database, max_attempts): + """ + rasie retryable error and exhause max_attempts + """ from google.api_core import exceptions + from google.cloud.firestore_v1.types import common + from google.cloud.firestore_v1.transaction import _EXCEED_ATTEMPTS_TEMPLATE - wrapped = _make__transactional(mock.sentinel.callable_) + to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) + wrapped = _make__transactional(to_wrap) - txn_id = b"failed" - transaction = _make_transaction_pb(txn_id, read_only=True, database=database) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. + txn_id = b"attempt_exhaustion" + transaction = _make_transaction_pb( + txn_id, database=database, max_attempts=max_attempts + ) - # Actually force the ``commit`` to fail (use ABORTED, but cannot - # retry since read-only). - exc = exceptions.Aborted("Read-only did a bad.") + # Actually force the ``commit`` to fail. + exc = exceptions.Aborted("Contention just once.") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc - with pytest.raises(exceptions.Aborted) as exc_info: - wrapped._maybe_commit(transaction) - assert exc_info.value is exc + # Call the __call__-able ``wrapped``. + with pytest.raises(ValueError) as exc_info: + wrapped(transaction, "here", there=1.5) - assert transaction._id == txn_id + err_msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) + assert exc_info.value.args == (err_msg,) + # should retain cause exception + assert exc_info.value.__cause__ == exc + + assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + assert to_wrap.call_count == max_attempts + to_wrap.assert_called_with(transaction, "here", there=1.5) + assert firestore_api.begin_transaction.call_count == max_attempts + options_ = common.TransactionOptions( + read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + ) + expected_calls = [ + mock.call( + request={ + "database": transaction._client._database_string, + "options": None if i == 0 else options_, + }, + metadata=transaction._client._rpc_metadata, + ) + for i in range(max_attempts) + ] + assert firestore_api.begin_transaction.call_args_list == expected_calls + assert firestore_api.commit.call_count == max_attempts + firestore_api.commit.assert_called_with( request={ "database": transaction._client._database_string, "writes": [], @@ -600,39 +614,9 @@ def test__transactional__maybe_commit_failure_read_only(database): }, metadata=transaction._client._rpc_metadata, ) - - -@pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__maybe_commit_failure_can_retry(database): - from google.api_core import exceptions - - wrapped = _make__transactional(mock.sentinel.callable_) - - txn_id = b"failed-but-retry" - transaction = _make_transaction_pb(txn_id, database=database) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. - - # Actually force the ``commit`` to fail. - exc = exceptions.Aborted("Read-write did a bad.") - firestore_api = transaction._client._firestore_api - firestore_api.commit.side_effect = exc - - succeeded = wrapped._maybe_commit(transaction) - assert not succeeded - - assert transaction._id == txn_id - assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id - - # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, - "writes": [], "transaction": txn_id, }, metadata=transaction._client._rpc_metadata, @@ -640,65 +624,55 @@ def test__transactional__maybe_commit_failure_can_retry(database): @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional__maybe_commit_failure_cannot_retry(database): +@pytest.mark.parametrize("max_attempts", [1, 5]) +def test_transactional___call__failure_readonly(database, max_attempts): + """ + readonly transaction should never retry + """ from google.api_core import exceptions + from google.cloud.firestore_v1.types import common - wrapped = _make__transactional(mock.sentinel.callable_) + to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) + wrapped = _make__transactional(to_wrap) - txn_id = b"failed-but-not-retryable" - transaction = _make_transaction_pb(txn_id, database=database) - transaction._id = txn_id # We won't call ``begin()``. - wrapped.current_id = txn_id # We won't call ``_pre_commit()``. - wrapped.retry_id = txn_id # We won't call ``_pre_commit()``. + txn_id = b"read_only_fail" + transaction = _make_transaction_pb( + txn_id, database=database, max_attempts=max_attempts, read_only=True + ) # Actually force the ``commit`` to fail. - exc = exceptions.InternalServerError("Real bad thing") + exc = exceptions.Aborted("Contention just once.") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc - with pytest.raises(exceptions.InternalServerError) as exc_info: - wrapped._maybe_commit(transaction) - assert exc_info.value is exc + # Call the __call__-able ``wrapped``. + with pytest.raises(exceptions.Aborted) as exc_info: + wrapped(transaction, "here", there=1.5) - assert transaction._id == txn_id + assert exc_info.value == exc + + assert transaction._id is None assert wrapped.current_id == txn_id assert wrapped.retry_id == txn_id # Verify mocks. - firestore_api.begin_transaction.assert_not_called() - firestore_api.rollback.assert_not_called() - firestore_api.commit.assert_called_once_with( + to_wrap.assert_called_once_with(transaction, "here", there=1.5) + firestore_api.begin_transaction.assert_called_once_with( request={ "database": transaction._client._database_string, - "writes": [], - "transaction": txn_id, + "options": common.TransactionOptions( + read_only=common.TransactionOptions.ReadOnly() + ), }, metadata=transaction._client._rpc_metadata, ) - - -@pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional___call__success_first_attempt(database): - to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) - wrapped = _make__transactional(to_wrap) - - txn_id = b"whole-enchilada" - transaction = _make_transaction_pb(txn_id, database=database) - result = wrapped(transaction, "a", b="c") - assert result is mock.sentinel.result - - assert transaction._id is None - assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id - - # Verify mocks. - to_wrap.assert_called_once_with(transaction, "a", b="c") - firestore_api = transaction._client._firestore_api - firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, + firestore_api.rollback.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "transaction": txn_id, + }, metadata=transaction._client._rpc_metadata, ) - firestore_api.rollback.assert_not_called() firestore_api.commit.assert_called_once_with( request={ "database": transaction._client._database_string, @@ -710,92 +684,102 @@ def test__transactional___call__success_first_attempt(database): @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional___call__success_second_attempt(database): +@pytest.mark.parametrize("max_attempts", [1, 5]) +def test_transactional___call__failure_with_non_retryable(database, max_attempts): + """ + call fails due to an exception that is not retryable. + Should rollback raise immediately + """ from google.api_core import exceptions - from google.cloud.firestore_v1.types import common - from google.cloud.firestore_v1.types import firestore - from google.cloud.firestore_v1.types import write to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) wrapped = _make__transactional(to_wrap) - txn_id = b"whole-enchilada" - transaction = _make_transaction_pb(txn_id, database=database) + txn_id = b"non_retryable" + transaction = _make_transaction_pb( + txn_id, database=database, max_attempts=max_attempts + ) - # Actually force the ``commit`` to fail on first / succeed on second. - exc = exceptions.Aborted("Contention junction.") + # Actually force the ``commit`` to fail. + exc = exceptions.InvalidArgument("non retryable") firestore_api = transaction._client._firestore_api - firestore_api.commit.side_effect = [ - exc, - firestore.CommitResponse(write_results=[write.WriteResult()]), - ] + firestore_api.commit.side_effect = exc # Call the __call__-able ``wrapped``. - result = wrapped(transaction, "a", b="c") - assert result is mock.sentinel.result + with pytest.raises(exceptions.InvalidArgument) as exc_info: + wrapped(transaction, "here", there=1.5) + + assert exc_info.value == exc assert transaction._id is None assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id # Verify mocks. - wrapped_call = mock.call(transaction, "a", b="c") - assert to_wrap.mock_calls, [wrapped_call == wrapped_call] - firestore_api = transaction._client._firestore_api - db_str = transaction._client._database_string - options_ = common.TransactionOptions( - read_write=common.TransactionOptions.ReadWrite(retry_transaction=txn_id) + to_wrap.assert_called_once_with(transaction, "here", there=1.5) + firestore_api.begin_transaction.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "options": None, + }, + metadata=transaction._client._rpc_metadata, ) - expected_calls = [ - mock.call( - request={"database": db_str, "options": None}, - metadata=transaction._client._rpc_metadata, - ), - mock.call( - request={"database": db_str, "options": options_}, - metadata=transaction._client._rpc_metadata, - ), - ] - assert firestore_api.begin_transaction.mock_calls == expected_calls - firestore_api.rollback.assert_not_called() - commit_call = mock.call( - request={"database": db_str, "writes": [], "transaction": txn_id}, + firestore_api.rollback.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "transaction": txn_id, + }, + metadata=transaction._client._rpc_metadata, + ) + firestore_api.commit.assert_called_once_with( + request={ + "database": transaction._client._database_string, + "writes": [], + "transaction": txn_id, + }, metadata=transaction._client._rpc_metadata, ) - assert firestore_api.commit.mock_calls == [commit_call, commit_call] @pytest.mark.parametrize("database", [None, "somedb"]) -def test__transactional___call__failure(database): +def test_transactional___call__failure_with_rollback_failure(database): + """ + Test second failure as part of rollback + should maintain first failure as __context__ + """ from google.api_core import exceptions - from google.cloud.firestore_v1.base_transaction import _EXCEED_ATTEMPTS_TEMPLATE to_wrap = mock.Mock(return_value=mock.sentinel.result, spec=[]) wrapped = _make__transactional(to_wrap) - txn_id = b"only-one-shot" - transaction = _make_transaction_pb(txn_id, max_attempts=1, database=database) + txn_id = b"non_retryable" + transaction = _make_transaction_pb(txn_id, database=database, max_attempts=1) # Actually force the ``commit`` to fail. - exc = exceptions.Aborted("Contention just once.") + exc = exceptions.InvalidArgument("first error") firestore_api = transaction._client._firestore_api firestore_api.commit.side_effect = exc + # also force a second error on rollback + rb_exc = exceptions.InternalServerError("second error") + firestore_api.rollback.side_effect = rb_exc # Call the __call__-able ``wrapped``. - with pytest.raises(ValueError) as exc_info: + # should raise second error with first error as __context__ + with pytest.raises(exceptions.InternalServerError) as exc_info: wrapped(transaction, "here", there=1.5) - err_msg = _EXCEED_ATTEMPTS_TEMPLATE.format(transaction._max_attempts) - assert exc_info.value.args == (err_msg,) + assert exc_info.value == rb_exc + assert exc_info.value.__context__ == exc assert transaction._id is None assert wrapped.current_id == txn_id - assert wrapped.retry_id == txn_id # Verify mocks. to_wrap.assert_called_once_with(transaction, "here", there=1.5) firestore_api.begin_transaction.assert_called_once_with( - request={"database": transaction._client._database_string, "options": None}, + request={ + "database": transaction._client._database_string, + "options": None, + }, metadata=transaction._client._rpc_metadata, ) firestore_api.rollback.assert_called_once_with(