fix: ensure transactions rollback on failure

google/cloud/firestore_v1/async_transaction.py

@@ -105,14 +105,18 @@ async def _begin(self, retry_id: bytes = None) -> None:
         )
         self._id = transaction_response.transaction
 
-    async def _rollback(self) -> None:
+    async def _rollback(self, source_exc=None) -> None:
         """Roll back the transaction.
 
+        Args:
+            source_exc (Optional[Exception]): The exception that caused the
+                rollback to occur. If an exception is created while rolling
+                back, it will be chained to this one.
         Raises:
             ValueError: If no transaction is in progress.
         """
         if not self.in_progress:
-            raise ValueError(_CANT_ROLLBACK)
+            raise ValueError(_CANT_ROLLBACK) from source_exc
 
         try:
             # NOTE: The response is just ``google.protobuf.Empty``.
@@ -123,6 +127,9 @@ async def _rollback(self) -> None:
                 },
                 metadata=self._client._rpc_metadata,
             )
+        except Exception as exc:  # pylint: disable=broad-except
+            # attach source_exc to the exception raised by rollback
+            raise exc from source_exc
         finally:
             self._clean_up()
 
@@ -223,10 +230,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 +253,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 +275,33 @@ 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 as exc:
+            await transaction._rollback(source_exc=exc)
+            raise exc
 
 
 def async_transactional(

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

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):
@@ -103,14 +103,18 @@ def _begin(self, retry_id: bytes = None) -> None:
         )
         self._id = transaction_response.transaction
 
-    def _rollback(self) -> None:
+    def _rollback(self, source_exc=None) -> None:
         """Roll back the transaction.
 
+        Args:
+            source_exc (Optional[Exception]): The exception that caused the
+                rollback to occur. If an exception is created while rolling
+                back, it will be chained to this one.
         Raises:
             ValueError: If no transaction is in progress.
         """
         if not self.in_progress:
-            raise ValueError(_CANT_ROLLBACK)
+            raise ValueError(_CANT_ROLLBACK) from source_exc
 
         try:
             # NOTE: The response is just ``google.protobuf.Empty``.
@@ -121,6 +125,9 @@ def _rollback(self) -> None:
                 },
                 metadata=self._client._rpc_metadata,
             )
+        except Exception as exc:
+            # attach source_exc to the exception raised by rollback
+            raise exc from source_exc
         finally:
             self._clean_up()
 
@@ -214,10 +221,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 +244,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 +266,33 @@ 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 as exc:  # noqa: B901
+            # rollback the transaction on any error
+            transaction._rollback(source_exc=exc)
+            raise exc
 
 
 def transactional(to_wrap: Callable) -> _Transactional:

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.