Expose transaction cancellation reasons (pynamodb#1144)

ikonst · web-flow · commit b03838cd2271 · 2023-01-14T19:09:23.000-05:00
This will allow the caller to know e.g. which of the transaction items ran into a conflict, which can sometimes allow app logic to better handle the error without needing to make more requests.
diff --git a/docs/api.rst b/docs/api.rst
@@ -44,3 +44,4 @@ Exceptions
 .. autoexception:: pynamodb.exceptions.InvalidStateError
 .. autoexception:: pynamodb.exceptions.AttributeDeserializationError
 .. autoexception:: pynamodb.exceptions.AttributeNullError
+.. autoclass:: pynamodb.exceptions.CancellationReason
diff --git a/docs/transaction.rst b/docs/transaction.rst
@@ -105,6 +105,10 @@ Now, say you make another attempt to debit one of the accounts when they don't h
         # Because the condition check on the account balance failed,
         # the entire transaction should be cancelled
         assert e.cause_response_code == 'TransactionCanceledException'
+        # the first 'update' was a reason for the cancellation
+        assert e.cancellation_reasons[0].code == 'ConditionalCheckFailed'
+        # the second 'update' wasn't a reason, but was cancelled too
+        assert e.cancellation_reasons[1] is None
 
         user1_statement.refresh()
         user2_statement.refresh()
diff --git a/pynamodb/connection/base.py b/pynamodb/connection/base.py
@@ -48,7 +48,8 @@
 from pynamodb.exceptions import (
     TableError, QueryError, PutError, DeleteError, UpdateError, GetError, ScanError, TableDoesNotExist,
     VerboseClientError,
-    TransactGetError, TransactWriteError)
+    TransactGetError, TransactWriteError, CancellationReason,
+)
 from pynamodb.expressions.condition import Condition
 from pynamodb.expressions.operand import Path
 from pynamodb.expressions.projection import create_projection_expression
@@ -465,7 +466,20 @@ def _make_api_call(self, operation_name: str, operation_kwargs: Dict, settings:
                     verbose_properties['table_name'] = operation_kwargs.get(TABLE_NAME)
 
                 try:
-                    raise VerboseClientError(botocore_expected_format, operation_name, verbose_properties)
+                    raise VerboseClientError(
+                        botocore_expected_format,
+                        operation_name,
+                        verbose_properties,
+                        cancellation_reasons=(
+                            (
+                                CancellationReason(
+                                    code=d['Code'],
+                                    message=d.get('Message'),
+                                ) if d['Code'] != 'None' else None
+                            )
+                            for d in data.get('CancellationReasons', [])
+                        ),
+                    )
                 except VerboseClientError as e:
                     if is_last_attempt_for_exceptions:
                         log.debug('Reached the maximum number of retry attempts: %s', attempt_number)
diff --git a/pynamodb/exceptions.py b/pynamodb/exceptions.py
@@ -1,8 +1,13 @@
 """
 PynamoDB exceptions
 """
-
-from typing import Any, Optional
+from dataclasses import dataclass
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing_extensions import Literal
 
 import botocore.exceptions
 
@@ -112,16 +117,63 @@ def __init__(self, table_name: str) -> None:
         super(TableDoesNotExist, self).__init__(msg)
 
 
+@dataclass
+class CancellationReason:
+    """
+    A reason for a transaction cancellation.
+    """
+    code: Literal[
+        'ConditionalCheckFailed',
+        'ItemCollectionSizeLimitExceeded',
+        'TransactionConflict',
+        'ProvisionedThroughputExceeded',
+        'ThrottlingError',
+        'ValidationError',
+    ]
+    message: Optional[str] = None
+
+
 class TransactWriteError(PynamoDBException):
     """
     Raised when a :class:`~pynamodb.transactions.TransactWrite` operation fails.
     """
 
+    @property
+    def cancellation_reasons(self) -> List[Optional[CancellationReason]]:
+        """
+        When :attr:`.cause_response_code` is ``TransactionCanceledException``, this property lists
+        cancellation reasons in the same order as the transaction items (one-to-one).
+        Items which were not part of the reason for cancellation would have :code:`None` as the value.
+
+        For a list of possible cancellation reasons and their semantics,
+        see `TransactWriteItems`_ in the AWS documentation.
+
+        .. _TransactWriteItems: https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_TransactWriteItems.html
+        """
+        if not isinstance(self.cause, VerboseClientError):
+            return []
+        return self.cause.cancellation_reasons
+
 
 class TransactGetError(PynamoDBException):
     """
     Raised when a :class:`~pynamodb.transactions.TransactGet` operation fails.
     """
+    @property
+    def cancellation_reasons(self) -> List[Optional[CancellationReason]]:
+        """
+        When :attr:`.cause_response_code` is ``TransactionCanceledException``, this property lists
+        cancellation reasons in the same order as the transaction items (one-to-one).
+        Items which were not part of the reason for cancellation would have :code:`None` as the value.
+
+        For a list of possible cancellation reasons and their semantics,
+        see `TransactGetItems`_ in the AWS documentation.
+
+        .. _TransactGetItems: https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_TransactGetItems.html
+        """
+        if not isinstance(self.cause, VerboseClientError):
+            return []
+        return self.cause.cancellation_reasons
 
 
 class InvalidStateError(PynamoDBException):
@@ -156,13 +208,23 @@ def prepend_path(self, attr_name: str) -> None:
 
 
 class VerboseClientError(botocore.exceptions.ClientError):
-    def __init__(self, error_response: Any, operation_name: str, verbose_properties: Optional[Any] = None) -> None:
+    def __init__(
+        self,
+        error_response: Dict[str, Any],
+        operation_name: str,
+        verbose_properties: Optional[Any] = None,
+        *,
+        cancellation_reasons: Iterable[Optional[CancellationReason]] = (),
+    ) -> None:
         """
         Like ClientError, but with a verbose message.
 
         :param error_response: Error response in shape expected by ClientError.
         :param operation_name: The name of the operation that failed.
         :param verbose_properties: A dict of properties to include in the verbose message.
+        :param cancellation_reasons: For `TransactionCanceledException` error code,
+          a list of cancellation reasons in the same order as the transaction's items (one to one).
+          For items which were not a reason for the transaction cancellation, :code:`None` will be the value.
         """
         if not verbose_properties:
             verbose_properties = {}
@@ -173,4 +235,9 @@ def __init__(self, error_response: Any, operation_name: str, verbose_properties:
             'operation: {{error_message}}'
         ).format(request_id=verbose_properties.get('request_id'), table_name=verbose_properties.get('table_name'))
 
-        super(VerboseClientError, self).__init__(error_response, operation_name)
+        self.cancellation_reasons = list(cancellation_reasons)
+
+        super(VerboseClientError, self).__init__(
+            error_response,  # type:ignore[arg-type]  # in stubs: botocore.exceptions._ClientErrorResponseTypeDef
+            operation_name,
+        )
diff --git a/tests/integration/test_transaction_integration.py b/tests/integration/test_transaction_integration.py
@@ -5,6 +5,7 @@
 import pytest
 
 from pynamodb.connection import Connection
+from pynamodb.exceptions import CancellationReason
 from pynamodb.exceptions import DoesNotExist, TransactWriteError, InvalidStateError
 
 
@@ -158,11 +159,32 @@ def test_transact_write__error__transaction_cancelled__condition_check_failure(c
             transaction.save(BankStatement(1), condition=(BankStatement.user_id.does_not_exist()))
     assert exc_info.value.cause_response_code == TRANSACTION_CANCELLED
     assert 'ConditionalCheckFailed' in exc_info.value.cause_response_message
+    assert exc_info.value.cancellation_reasons == [
+        CancellationReason(code='ConditionalCheckFailed', message='The conditional request failed'),
+        CancellationReason(code='ConditionalCheckFailed', message='The conditional request failed'),
+    ]
     assert isinstance(exc_info.value.cause, botocore.exceptions.ClientError)
     assert User.Meta.table_name in exc_info.value.cause.MSG_TEMPLATE
     assert BankStatement.Meta.table_name in exc_info.value.cause.MSG_TEMPLATE
 
 
+@pytest.mark.ddblocal
+def test_transact_write__error__transaction_cancelled__partial_failure(connection):
+    User(2).delete()
+    BankStatement(2).save()
+
+    # attempt to do this as a transaction with the condition that they don't already exist
+    with pytest.raises(TransactWriteError) as exc_info:
+        with TransactWrite(connection=connection) as transaction:
+            transaction.save(User(2), condition=(User.user_id.does_not_exist()))
+            transaction.save(BankStatement(2), condition=(BankStatement.user_id.does_not_exist()))
+    assert exc_info.value.cause_response_code == TRANSACTION_CANCELLED
+    assert exc_info.value.cancellation_reasons == [
+        None,
+        CancellationReason(code='ConditionalCheckFailed', message='The conditional request failed'),
+    ]
+
+
 @pytest.mark.ddblocal
 def test_transact_write__error__multiple_operations_on_same_record(connection):
     BankStatement(1).save()
