consolidate crypto config building logic across encrypted Client/Resource/Table

Author: mattsb42-aws

src/dynamodb_encryption_sdk/encrypted/client.py

@@ -17,20 +17,20 @@
 import botocore.client
 
 from dynamodb_encryption_sdk.internal.utils import (
+    crypto_config_from_cache, crypto_config_from_kwargs,
     decrypt_batch_get_item, decrypt_get_item, decrypt_multi_get,
     encrypt_batch_write_item, encrypt_put_item, TableInfoCache
 )
 from dynamodb_encryption_sdk.material_providers import CryptographicMaterialsProvider
-from dynamodb_encryption_sdk.structures import AttributeActions, EncryptionContext
-from . import CryptoConfig
+from dynamodb_encryption_sdk.structures import AttributeActions
 from .item import decrypt_dynamodb_item, encrypt_dynamodb_item
 
 __all__ = ('EncryptedClient',)
 
 
 @attr.s
 class EncryptedClient(object):
-    # pylint: disable=too-few-public-methods
+    # pylint: disable=too-few-public-methods,too-many-instance-attributes
     """High-level helper class to provide a familiar interface to encrypted tables.
 
     .. note::
@@ -68,40 +68,50 @@ def __attrs_post_init__(self):
             client=self._client,
             auto_refresh_table_indexes=self._auto_refresh_table_indexes
         )
+        self._table_crypto_config = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
+            crypto_config_from_cache,
+            materials_provider=self._materials_provider,
+            attribute_actions=self._attribute_actions,
+            table_info_cache=self._table_info_cache
+        )
+        self._item_crypto_config = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
+            crypto_config_from_kwargs,
+            fallback=self._table_crypto_config
+        )
         self.get_item = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             decrypt_get_item,
             decrypt_method=decrypt_dynamodb_item,
-            crypto_config_method=self._table_crypto_config,
+            crypto_config_method=self._item_crypto_config,
             read_method=self._client.get_item
         )
         self.put_item = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             encrypt_put_item,
             encrypt_method=encrypt_dynamodb_item,
-            crypto_config_method=self._table_crypto_config,
+            crypto_config_method=self._item_crypto_config,
             write_method=self._client.put_item
         )
         self.query = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             decrypt_multi_get,
             decrypt_method=decrypt_dynamodb_item,
-            crypto_config_method=self._table_crypto_config,
+            crypto_config_method=self._item_crypto_config,
             read_method=self._client.query
         )
         self.scan = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             decrypt_multi_get,
             decrypt_method=decrypt_dynamodb_item,
-            crypto_config_method=self._table_crypto_config,
+            crypto_config_method=self._item_crypto_config,
             read_method=self._client.scan
         )
         self.batch_get_item = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             decrypt_batch_get_item,
             decrypt_method=decrypt_dynamodb_item,
-            crypto_config_method=self._batch_crypto_config,
+            crypto_config_method=self._table_crypto_config,
             read_method=self._client.batch_get_item
         )
         self.batch_write_item = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             encrypt_batch_write_item,
             encrypt_method=encrypt_dynamodb_item,
-            crypto_config_method=self._batch_crypto_config,
+            crypto_config_method=self._table_crypto_config,
             write_method=self._client.batch_write_item
         )
 
@@ -115,47 +125,6 @@ def __getattr__(self, name):
         """
         return getattr(self._client, name)
 
-    def _crypto_config(self, table_name, **kwargs):
-        """Pull all encryption-specific parameters from the request and use them to build a crypto config.
-
-        :returns: crypto config and updated kwargs
-        :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
-        """
-        crypto_config = kwargs.pop('crypto_config', None)
-
-        if crypto_config is not None:
-            return crypto_config, kwargs
-
-        table_info = self._table_info_cache.table_info(table_name)
-
-        attribute_actions = self._attribute_actions.copy()
-        attribute_actions.set_index_keys(*table_info.protected_index_keys())
-
-        crypto_config = CryptoConfig(
-            materials_provider=self._materials_provider,
-            encryption_context=EncryptionContext(**table_info.encryption_context_values),
-            attribute_actions=attribute_actions
-        )
-        return crypto_config, kwargs
-
-    def _table_crypto_config(self, **kwargs):
-        """Pull all encryption-specific parameters from the request and use them to build
-        a crypto config for a single-table operation.
-
-        :returns: crypto config and updated kwargs
-        :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
-        """
-        return self._crypto_config(kwargs['TableName'], **kwargs)
-
-    def _batch_crypto_config(self, table_name):
-        """Build a crypto config for a specific table.
-
-        :param str table_name: Table for which to build crypto config
-        :returns: crypto config
-        :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig
-        """
-        return self._crypto_config(table_name)[0]
-
     def update_item(self, **kwargs):
         """Update item is not yet supported."""
         raise NotImplementedError('"update_item" is not yet implemented')

src/dynamodb_encryption_sdk/encrypted/resource.py

@@ -17,10 +17,11 @@
 from boto3.resources.base import ServiceResource
 from boto3.resources.collection import CollectionManager
 
-from dynamodb_encryption_sdk.internal.utils import decrypt_batch_get_item, encrypt_batch_write_item, TableInfoCache
+from dynamodb_encryption_sdk.internal.utils import (
+    crypto_config_from_cache, decrypt_batch_get_item, encrypt_batch_write_item, TableInfoCache
+)
 from dynamodb_encryption_sdk.material_providers import CryptographicMaterialsProvider
-from dynamodb_encryption_sdk.structures import AttributeActions, EncryptionContext
-from . import CryptoConfig
+from dynamodb_encryption_sdk.structures import AttributeActions
 from .item import decrypt_python_item, encrypt_python_item
 from .table import EncryptedTable
 
@@ -137,6 +138,12 @@ def __attrs_post_init__(self):
             client=self._resource.meta.client,
             auto_refresh_table_indexes=self._auto_refresh_table_indexes
         )
+        self._crypto_config = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
+            crypto_config_from_cache,
+            materials_provider=self._materials_provider,
+            attribute_actions=self._attribute_actions,
+            table_info_cache=self._table_info_cache
+        )
         self.tables = EncryptedTablesCollectionManager(  # attrs confuses pylint: disable=attribute-defined-outside-init
             collection=self._resource.tables,
             materials_provider=self._materials_provider,
@@ -166,24 +173,6 @@ def __getattr__(self, name):
         """
         return getattr(self._resource, name)
 
-    def _crypto_config(self, table_name):
-        """Pull all encryption-specific parameters from the request and use them to build a crypto config.
-
-        :returns: crypto config
-        :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig
-        """
-        table_info = self._table_info_cache.table_info(table_name)
-
-        attribute_actions = self._attribute_actions.copy()
-        attribute_actions.set_index_keys(*table_info.protected_index_keys())
-
-        crypto_config = CryptoConfig(
-            materials_provider=self._materials_provider,
-            encryption_context=EncryptionContext(**table_info.encryption_context_values),
-            attribute_actions=attribute_actions
-        )
-        return crypto_config
-
     def Table(self, name, **kwargs):
         # naming chosen to align with boto3 resource name, so pylint: disable=invalid-name
         """Creates an EncryptedTable resource.

src/dynamodb_encryption_sdk/encrypted/table.py

@@ -16,10 +16,12 @@
 import attr
 from boto3.resources.base import ServiceResource
 
-from dynamodb_encryption_sdk.internal.utils import decrypt_get_item, decrypt_multi_get, encrypt_put_item
+from dynamodb_encryption_sdk.internal.utils import (
+    crypto_config_from_kwargs, crypto_config_from_table_info,
+    decrypt_get_item, decrypt_multi_get, encrypt_put_item
+)
 from dynamodb_encryption_sdk.material_providers import CryptographicMaterialsProvider
-from dynamodb_encryption_sdk.structures import AttributeActions, EncryptionContext, TableInfo
-from . import CryptoConfig
+from dynamodb_encryption_sdk.structures import AttributeActions, TableInfo
 from .item import decrypt_python_item, encrypt_python_item
 
 __all__ = ('EncryptedTable',)
@@ -84,6 +86,15 @@ def __attrs_post_init__(self):
         self._attribute_actions = self._attribute_actions.copy()
         self._attribute_actions.set_index_keys(*self._table_info.protected_index_keys())
 
+        self._crypto_config = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
+            crypto_config_from_kwargs,
+            fallback=partial(
+                crypto_config_from_table_info,
+                materials_provider=self._materials_provider,
+                attribute_actions=self._attribute_actions,
+                table_info=self._table_info
+            )
+        )
         self.get_item = partial(  # attrs confuses pylint: disable=attribute-defined-outside-init
             decrypt_get_item,
             decrypt_method=decrypt_python_item,
@@ -122,21 +133,3 @@ def __getattr__(self, name):
     def update_item(self, **kwargs):
         """Update item is not yet supported."""
         raise NotImplementedError('"update_item" is not yet implemented')
-
-    def _crypto_config(self, **kwargs):
-        """Pull all encryption-specific parameters from the request and use them to build a crypto config.
-
-        :returns: crypto config and updated kwargs
-        :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
-        """
-        crypto_config = kwargs.pop('crypto_config', None)
-
-        if crypto_config is not None:
-            return crypto_config, kwargs
-
-        crypto_config = CryptoConfig(
-            materials_provider=self._materials_provider,
-            encryption_context=EncryptionContext(**self._table_info.encryption_context_values),
-            attribute_actions=self._attribute_actions
-        )
-        return crypto_config, kwargs

src/dynamodb_encryption_sdk/internal/utils.py

@@ -14,9 +14,10 @@
 import attr
 import botocore.client
 
+from dynamodb_encryption_sdk.encrypted import CryptoConfig
 from dynamodb_encryption_sdk.exceptions import InvalidArgumentError
 from dynamodb_encryption_sdk.internal.str_ops import to_bytes
-from dynamodb_encryption_sdk.structures import TableInfo
+from dynamodb_encryption_sdk.structures import EncryptionContext, TableInfo
 
 try:  # Python 3.5.0 and 3.5.1 have incompatible typing modules
     from typing import Callable, Dict, Text  # noqa pylint: disable=unused-import
@@ -26,6 +27,7 @@
 
 __all__ = (
     'sorted_key_map', 'TableInfoCache',
+    'crypto_config_from_kwargs', 'crypto_config_from_table_info', 'crypto_config_from_cache',
     'decrypt_get_item', 'decrypt_multi_get', 'decrypt_batch_get_item',
     'encrypt_put_item', 'encrypt_batch_write_item'
 )
@@ -97,6 +99,50 @@ def validate_get_arguments(kwargs):
         raise InvalidArgumentError('Scan "Select" value of "{}" is not supported'.format(kwargs['Select']))
 
 
+def crypto_config_from_kwargs(fallback, **kwargs):
+    """Pull all encryption-specific parameters from the request and use them to build a crypto config.
+
+    :returns: crypto config and updated kwargs
+    :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
+    """
+    try:
+        crypto_config = kwargs.pop('crypto_config')
+    except KeyError:
+        try:
+            fallback_kwargs = {'table_name': kwargs['TableName']}
+        except KeyError:
+            fallback_kwargs = {}
+        crypto_config = fallback(**fallback_kwargs)
+    return crypto_config, kwargs
+
+
+def crypto_config_from_table_info(materials_provider, attribute_actions, table_info):
+    """Build a crypto config from the provided values and table info.
+
+    :returns: crypto config and updated kwargs
+    :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
+    """
+    return CryptoConfig(
+        materials_provider=materials_provider,
+        encryption_context=EncryptionContext(**table_info.encryption_context_values),
+        attribute_actions=attribute_actions
+    )
+
+
+def crypto_config_from_cache(materials_provider, attribute_actions, table_info_cache, table_name):
+    """Build a crypto config from the provided values, loading the table info from the provided cache.
+
+    :returns: crypto config and updated kwargs
+    :rtype: dynamodb_encryption_sdk.encrypted.CryptoConfig and dict
+    """
+    table_info = table_info_cache.table_info(table_name)
+
+    attribute_actions = attribute_actions.copy()
+    attribute_actions.set_index_keys(*table_info.protected_index_keys())
+
+    return crypto_config_from_table_info(materials_provider, attribute_actions, table_info)
+
+
 def decrypt_multi_get(decrypt_method, crypto_config_method, read_method, **kwargs):
     # type: (Callable, Callable, Callable, **Any) -> Dict
     # TODO: narrow this down
@@ -156,7 +202,7 @@ def decrypt_batch_get_item(decrypt_method, crypto_config_method, read_method, **
         if request_crypto_config is not None:
             crypto_config = request_crypto_config
         else:
-            crypto_config = crypto_config_method(table_name)
+            crypto_config = crypto_config_method(table_name=table_name)
 
         for pos, value in enumerate(items):
             items[pos] = decrypt_method(
@@ -198,7 +244,7 @@ def encrypt_batch_write_item(encrypt_method, crypto_config_method, write_method,
         if request_crypto_config is not None:
             crypto_config = request_crypto_config
         else:
-            crypto_config = crypto_config_method(table_name)
+            crypto_config = crypto_config_method(table_name=table_name)
 
         for pos, value in enumerate(items):
             for request_type, item in value.items():