Merge branch 'gh-170'

Author: ojii

src/aiodynamo/errors.py

@@ -1,5 +1,12 @@
 import json
-from typing import Any, Dict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+
+@dataclass(frozen=True)
+class CancellationReason:
+    code: str
+    message: str
 
 
 class AIODynamoError(Exception):
@@ -110,7 +117,17 @@ class PointInTimeRecoveryUnavailable(AIODynamoError):
 
 
 class TransactionCanceled(AIODynamoError):
-    pass
+    cancellation_reasons: List[Union[CancellationReason, None]]
+
+    def __init__(self, body: Dict[str, Any]):
+        self.body = body
+        self.cancellation_reasons: List[CancellationReason] = [
+            CancellationReason(reason["Code"], reason["Message"])
+            if reason["Code"] != "None"
+            else None
+            for reason in self.body["CancellationReasons"]
+        ]
+        super().__init__(body)
 
 
 class TransactionEmpty(AIODynamoError):
@@ -207,17 +224,7 @@ def exception_from_response(status: int, body: bytes) -> Exception:
         return ServiceUnavailable()
     try:
         data = json.loads(body)
-        error = ERRORS[data["__type"].split("#", 1)[-1]](data)
-        if isinstance(error, TransactionCanceled):
-            error = extract_error_from_transaction_canceled(data)
+        error: Exception = ERRORS[data["__type"].split("#", 1)[-1]](data)
         return error
     except Exception:
         return UnknownError(status, body)
-
-
-def extract_error_from_transaction_canceled(data: Dict[str, Any]) -> AIODynamoError:
-    try:
-        error = data["CancellationReasons"][0]
-        return ERRORS[f"{error['Code']}Exception"](error["Message"])
-    except Exception:
-        return ERRORS[data["__type"].split("#", 1)[-1]](data)

tests/integration/test_client.py

@@ -624,14 +624,36 @@ async def test_transact_write_items_put(client: Client, table: TableName) -> Non
     await client.transact_write_items(items=puts)
     assert len([item async for item in client.query(table, HashKey("h", "h"))]) == 2
 
-    with pytest.raises(errors.ConditionalCheckFailed):
+    with pytest.raises(errors.TransactionCanceled) as excinfo:
         put = Put(
             table=table,
             item={"h": "h", "r": "0", "s": "initial"},
             condition=F("h").does_not_exist(),
         )
         await client.transact_write_items(items=[put])
 
+    assert len(excinfo.value.cancellation_reasons) == 1
+    assert excinfo.value.cancellation_reasons[0] is not None
+    assert excinfo.value.cancellation_reasons[0].code == "ConditionalCheckFailed"
+
+    with pytest.raises(errors.TransactionCanceled) as excinfo:
+        put = Put(
+            table=table,
+            item={"h": "h", "r": "3", "s": "initial"},
+            condition=F("h").does_not_exist(),
+        )
+        put_fail = Put(
+            table=table,
+            item={"h": "h", "r": "0", "s": "initial"},
+            condition=F("h").does_not_exist(),
+        )
+        await client.transact_write_items(items=[put, put_fail])
+
+    assert len(excinfo.value.cancellation_reasons) == 2
+    assert excinfo.value.cancellation_reasons[0] is None
+    assert excinfo.value.cancellation_reasons[1] is not None
+    assert excinfo.value.cancellation_reasons[1].code == "ConditionalCheckFailed"
+
 
 @pytest.mark.usefixtures("supports_transactions")
 async def test_transact_write_items_update(client: Client, table: TableName) -> None:
@@ -647,7 +669,7 @@ async def test_transact_write_items_update(client: Client, table: TableName) ->
     query = await client.query_single_page(table, HashKey("h", "h"))
     assert query.items[0]["s"] == "changed"
 
-    with pytest.raises(errors.ConditionalCheckFailed):
+    with pytest.raises(errors.TransactionCanceled) as excinfo:
         update = Update(
             table=table,
             key={"h": "h", "r": "1"},
@@ -656,6 +678,10 @@ async def test_transact_write_items_update(client: Client, table: TableName) ->
         )
         await client.transact_write_items(items=[update])
 
+    assert len(excinfo.value.cancellation_reasons) == 1
+    assert excinfo.value.cancellation_reasons[0] is not None
+    assert excinfo.value.cancellation_reasons[0].code == "ConditionalCheckFailed"
+
 
 @pytest.mark.usefixtures("supports_transactions")
 async def test_transact_write_items_delete(client: Client, table: TableName) -> None:
@@ -670,14 +696,19 @@ async def test_transact_write_items_delete(client: Client, table: TableName) ->
     assert len([item async for item in client.query(table, HashKey("h", "h"))]) == 0
 
     await client.put_item(table=table, item={"h": "h", "r": "1", "s": "initial"})
-    with pytest.raises(errors.ConditionalCheckFailed):
+
+    with pytest.raises(errors.TransactionCanceled) as excinfo:
         delete = Delete(
             table=table,
             key={"h": "h", "r": "1"},
             condition=F("s").not_equals("initial"),
         )
         await client.transact_write_items(items=[delete])
 
+    assert len(excinfo.value.cancellation_reasons) == 1
+    assert excinfo.value.cancellation_reasons[0] is not None
+    assert excinfo.value.cancellation_reasons[0].code == "ConditionalCheckFailed"
+
 
 @pytest.mark.usefixtures("supports_transactions")
 async def test_transact_write_items_condition_check(
@@ -687,9 +718,13 @@ async def test_transact_write_items_condition_check(
     condition = ConditionCheck(
         table=table, key={"h": "h", "r": "1"}, condition=F("s").not_equals("initial")
     )
-    with pytest.raises(errors.ConditionalCheckFailed):
+    with pytest.raises(errors.TransactionCanceled) as excinfo:
         await client.transact_write_items(items=[condition])
 
+    assert len(excinfo.value.cancellation_reasons) == 1
+    assert excinfo.value.cancellation_reasons[0] is not None
+    assert excinfo.value.cancellation_reasons[0].code == "ConditionalCheckFailed"
+
     condition = ConditionCheck(
         table=table, key={"h": "h", "r": "1"}, condition=F("s").equals("initial")
     )