Handle unprocessed items in batch write responses

src/dynamodb_encryption_sdk/internal/utils.py

@@ -16,17 +16,20 @@
     No guarantee is provided on the modules and APIs within this
     namespace staying consistent. Directly reference at your own risk.
 """
+import copy
+from functools import partial
+
 import attr
 import botocore.client
 
 from dynamodb_encryption_sdk.encrypted import CryptoConfig
 from dynamodb_encryption_sdk.encrypted.item import decrypt_python_item, encrypt_python_item
 from dynamodb_encryption_sdk.exceptions import InvalidArgumentError
-from dynamodb_encryption_sdk.structures import EncryptionContext, TableInfo
+from dynamodb_encryption_sdk.structures import CryptoAction, EncryptionContext, TableInfo
 from dynamodb_encryption_sdk.transform import dict_to_ddb
 
 try:  # Python 3.5.0 and 3.5.1 have incompatible typing modules
-    from typing import Any, Callable, Dict, Text  # noqa pylint: disable=unused-import
+    from typing import Any, Bool, Callable, Dict, Text  # noqa pylint: disable=unused-import
 except ImportError:  # pragma: no cover
     # We only actually need these imports when running the mypy checks
     pass
@@ -271,19 +274,22 @@ def encrypt_batch_write_item(encrypt_method, crypto_config_method, write_method,
     """Transparently encrypt multiple items before putting them in a batch request.
 
     :param callable encrypt_method: Method to use to encrypt items
-    :param callable crypto_config_method: Method that accepts ``kwargs`` and provides a :class:`CryptoConfig`
+    :param callable crypto_config_method: Method that accepts a table name string and provides a :class:`CryptoConfig`
     :param callable write_method: Method that writes to the table
     :param **kwargs: Keyword arguments to pass to ``write_method``
     :return: DynamoDB response
     :rtype: dict
     """
     request_crypto_config = kwargs.pop("crypto_config", None)
+    table_crypto_configs = {}
+    plaintext_items = copy.deepcopy(kwargs["RequestItems"])
 
     for table_name, items in kwargs["RequestItems"].items():
         if request_crypto_config is not None:
             crypto_config = request_crypto_config
         else:
             crypto_config = crypto_config_method(table_name=table_name)
+        table_crypto_configs[table_name] = crypto_config
 
         for pos, value in enumerate(items):
             for request_type, item in value.items():
@@ -293,4 +299,91 @@ def encrypt_batch_write_item(encrypt_method, crypto_config_method, write_method,
                         item=item["Item"],
                         crypto_config=crypto_config.with_item(_item_transformer(encrypt_method)(item["Item"])),
                     )
-    return write_method(**kwargs)
+
+    response = write_method(**kwargs)
+    return _process_batch_write_response(plaintext_items, response, table_crypto_configs)
+
+
+def _process_batch_write_response(request, response, table_crypto_config):
+    # type: (Dict, Dict, Dict[Text, CryptoConfig]) -> Dict
+    """Handle unprocessed items in the response from a transparently encrypted write.
+
+    :param dict request: The DynamoDB plaintext request dictionary
+    :param dict response: The DynamoDB response from the batch operation
+    :param Dict[Text, CryptoConfig] table_crypto_config: table level CryptoConfig used in encrypting the request items
+    :return: DynamoDB response, with any unprocessed items reverted back to the original plaintext values
+    :rtype: dict
+    """
+    try:
+        unprocessed_items = response["UnprocessedItems"]
+    except KeyError:
+        return response
+
+    # Unprocessed items need to be returned in their original state
+    for table_name, unprocessed in unprocessed_items.items():
+        original_items = request[table_name]
+        crypto_config = table_crypto_config[table_name]
+
+        if crypto_config.encryption_context.partition_key_name:
+            items_match = partial(_item_keys_match, crypto_config)
+        else:
+            items_match = partial(_item_attributes_match, crypto_config)
+
+        for pos, operation in enumerate(unprocessed):
+            for request_type, item in operation.items():
+                if request_type != "PutRequest":
+                    continue
+
+                for plaintext_item in original_items:
+                    if plaintext_item.get(request_type) and items_match(
+                        plaintext_item[request_type]["Item"], item["Item"]
+                    ):
+                        unprocessed[pos] = plaintext_item.copy()
+                        break
+
+    return response
+
+
+def _item_keys_match(crypto_config, item1, item2):
+    # type: (CryptoConfig, Dict, Dict) -> Bool
+    """Determines whether the values in the primary and sort keys (if they exist) are the same
+
+    :param CryptoConfig crypto_config: CryptoConfig used in encrypting the given items
+    :param dict item1: The first item to compare
+    :param dict item2: The second item to compare
+    :return: Bool response, True if the key attributes match
+    :rtype: bool
+    """
+    partition_key_name = crypto_config.encryption_context.partition_key_name
+    sort_key_name = crypto_config.encryption_context.sort_key_name
+
+    partition_keys_match = item1[partition_key_name] == item2[partition_key_name]
+
+    if sort_key_name is None:
+        return partition_keys_match
+
+    return partition_keys_match and item1[sort_key_name] == item2[sort_key_name]
+
+
+def _item_attributes_match(crypto_config, plaintext_item, encrypted_item):
+    # type: (CryptoConfig, Dict, Dict) -> Bool
+    """Determines whether the unencrypted values in the plaintext items attributes are the same as those in the
+    encrypted item. Essentially this uses brute force to cover when we don't know the primary and sort
+    index attribute names, since they can't be encrypted.
+
+    :param CryptoConfig crypto_config: CryptoConfig used in encrypting the given items
+    :param dict plaintext_item: The plaintext item
+    :param dict encrypted_item: The encrypted item
+    :return: Bool response, True if the unencrypted attributes in the plaintext item match those in
+    the encrypted item
+    :rtype: bool
+    """
+
+    for name, value in plaintext_item.items():
+        if crypto_config.attribute_actions.action(name) == CryptoAction.ENCRYPT_AND_SIGN:
+            continue
+
+        if encrypted_item.get(name) != value:
+            return False
+
+    return True

test/functional/encrypted/test_client.py

@@ -17,6 +17,8 @@
 from ..functional_test_utils import example_table  # noqa pylint: disable=unused-import
 from ..functional_test_utils import (
     TEST_TABLE_NAME,
+    build_static_jce_cmp,
+    client_batch_items_unprocessed_check,
     client_cycle_batch_items_check,
     client_cycle_batch_items_check_paginators,
     client_cycle_single_item_check,
@@ -53,6 +55,12 @@ def _client_cycle_batch_items_check_paginators(materials_provider, initial_actio
     )
 
 
+def _client_batch_items_unprocessed_check(materials_provider, initial_actions, initial_item):
+    client_batch_items_unprocessed_check(
+        materials_provider, initial_actions, initial_item, TEST_TABLE_NAME, "us-west-2"
+    )
+
+
 def test_ephemeral_item_cycle(example_table, some_cmps, parametrized_actions, parametrized_item):
     """Test a small number of curated CMPs against a small number of curated items."""
     _client_cycle_single_item_check(some_cmps, parametrized_actions, parametrized_item)
@@ -68,6 +76,12 @@ def test_ephemeral_batch_item_cycle_paginators(example_table, some_cmps, paramet
     _client_cycle_batch_items_check_paginators(some_cmps, parametrized_actions, parametrized_item)
 
 
+def test_batch_item_unprocessed(example_table, parametrized_actions, parametrized_item):
+    """Test Unprocessed Items handling with a single ephemeral static CMP against a small number of curated items."""
+    cmp = build_static_jce_cmp("AES", 256, "HmacSHA256", 256)
+    _client_batch_items_unprocessed_check(cmp, parametrized_actions, parametrized_item)
+
+
 @pytest.mark.slow
 def test_ephemeral_item_cycle_slow(example_table, all_the_cmps, parametrized_actions, parametrized_item):
     """Test ALL THE CMPS against a small number of curated items."""

test/functional/encrypted/test_resource.py

@@ -13,9 +13,11 @@
 """Functional tests for ``dynamodb_encryption_sdk.encrypted.resource``."""
 import pytest
 
-from ..functional_test_utils import example_table  # noqa pylint: disable=unused-import
+from ..functional_test_utils import example_table # noqa pylint: disable=unused-import
 from ..functional_test_utils import (
     TEST_TABLE_NAME,
+    build_static_jce_cmp,
+    resource_batch_items_unprocessed_check,
     resource_cycle_batch_items_check,
     set_parametrized_actions,
     set_parametrized_cmp,
@@ -35,11 +37,24 @@ def _resource_cycle_batch_items_check(materials_provider, initial_actions, initi
     resource_cycle_batch_items_check(materials_provider, initial_actions, initial_item, TEST_TABLE_NAME, "us-west-2")
 
 
+def _resource_batch_items_unprocessed_check(materials_provider, initial_actions, initial_item):
+    resource_batch_items_unprocessed_check(
+        materials_provider, initial_actions, initial_item, TEST_TABLE_NAME, "us-west-2"
+    )
+
+
 def test_ephemeral_batch_item_cycle(example_table, some_cmps, parametrized_actions, parametrized_item):
     """Test a small number of curated CMPs against a small number of curated items."""
     _resource_cycle_batch_items_check(some_cmps, parametrized_actions, parametrized_item)
 
 
+def test_batch_item_unprocessed(example_table, parametrized_actions, parametrized_item):
+    """Test Unprocessed Items handling with a single ephemeral static CMP against a small number of curated items."""
+    _resource_batch_items_unprocessed_check(
+        build_static_jce_cmp("AES", 256, "HmacSHA256", 256), parametrized_actions, parametrized_item
+    )
+
+
 @pytest.mark.travis_isolation
 @pytest.mark.slow
 def test_ephemeral_batch_item_cycle_slow(example_table, all_the_cmps, parametrized_actions, parametrized_item):

test/functional/encrypted/test_table.py

@@ -14,12 +14,14 @@
 import hypothesis
 import pytest
 
-from ..functional_test_utils import example_table  # noqa pylint: disable=unused-import
+from ..functional_test_utils import example_table # noqa pylint: disable=unused-import
 from ..functional_test_utils import (
     TEST_TABLE_NAME,
+    build_static_jce_cmp,
     set_parametrized_actions,
     set_parametrized_cmp,
     set_parametrized_item,
+    table_batch_writer_unprocessed_items_check,
     table_cycle_batch_writer_check,
     table_cycle_check,
 )
@@ -48,6 +50,14 @@ def test_ephemeral_item_cycle_batch_writer(example_table, some_cmps, parametrize
     table_cycle_batch_writer_check(some_cmps, parametrized_actions, parametrized_item, TEST_TABLE_NAME, "us-west-2")
 
 
+def test_batch_writer_unprocessed(example_table, parametrized_actions, parametrized_item):
+    """Test Unprocessed Items handling with a single ephemeral static CMP against a small number of curated items."""
+    cmp = build_static_jce_cmp("AES", 256, "HmacSHA256", 256)
+    table_batch_writer_unprocessed_items_check(
+        cmp, parametrized_actions, parametrized_item, TEST_TABLE_NAME, "us-west-2"
+    )
+
+
 @pytest.mark.slow
 def test_ephemeral_item_cycle_slow(example_table, all_the_cmps, parametrized_actions, parametrized_item):
     """Test ALL THE CMPS against a small number of curated items."""