Replace AttributesToGet with ProjectionExpression / ExpressionAttributeNames (#319)

jpinner-lyft · web-flow · commit 11650c9f65e2 · 2017-08-02T11:24:11.000-07:00
diff --git a/pynamodb/connection/base.py b/pynamodb/connection/base.py
@@ -6,6 +6,7 @@
 import logging
 import math
 import random
+import re
 import time
 import uuid
 from base64 import b64decode
@@ -26,7 +27,7 @@
     COMPARISON_OPERATOR, EXCLUSIVE_START_KEY, SCAN_INDEX_FORWARD, SCAN_FILTER_VALUES, ATTR_DEFINITIONS,
     BATCH_WRITE_ITEM, CONSISTENT_READ, ATTR_VALUE_LIST, DESCRIBE_TABLE, KEY_CONDITIONS,
     BATCH_GET_ITEM, DELETE_REQUEST, SELECT_VALUES, RETURN_VALUES, REQUEST_ITEMS, ATTR_UPDATES,
-    ATTRS_TO_GET, SERVICE_NAME, DELETE_ITEM, PUT_REQUEST, UPDATE_ITEM, SCAN_FILTER, TABLE_NAME,
+    PROJECTION_EXPRESSION, SERVICE_NAME, DELETE_ITEM, PUT_REQUEST, UPDATE_ITEM, SCAN_FILTER, TABLE_NAME,
     INDEX_NAME, KEY_SCHEMA, ATTR_NAME, ATTR_TYPE, TABLE_KEY, EXPECTED, KEY_TYPE, GET_ITEM, UPDATE,
     PUT_ITEM, SELECT, ACTION, EXISTS, VALUE, LIMIT, QUERY, SCAN, ITEM, LOCAL_SECONDARY_INDEXES,
     KEYS, KEY, EQ, SEGMENT, TOTAL_SEGMENTS, CREATE_TABLE, PROVISIONED_THROUGHPUT, READ_CAPACITY_UNITS,
@@ -35,7 +36,7 @@
     CONSUMED_CAPACITY, CAPACITY_UNITS, QUERY_FILTER, QUERY_FILTER_VALUES, CONDITIONAL_OPERATOR,
     CONDITIONAL_OPERATORS, NULL, NOT_NULL, SHORT_ATTR_TYPES, DELETE,
     ITEMS, DEFAULT_ENCODING, BINARY_SHORT, BINARY_SET_SHORT, LAST_EVALUATED_KEY, RESPONSES, UNPROCESSED_KEYS,
-    UNPROCESSED_ITEMS, STREAM_SPECIFICATION, STREAM_VIEW_TYPE, STREAM_ENABLED)
+    UNPROCESSED_ITEMS, STREAM_SPECIFICATION, STREAM_VIEW_TYPE, STREAM_ENABLED, EXPRESSION_ATTRIBUTE_NAMES)
 from pynamodb.exceptions import (
     TableError, QueryError, PutError, DeleteError, UpdateError, GetError, ScanError, TableDoesNotExist,
     VerboseClientError
@@ -45,6 +46,7 @@
 from pynamodb.types import HASH, RANGE
 
 BOTOCORE_EXCEPTIONS = (BotoCoreError, ClientError)
+PATH_SEGMENT_REGEX = re.compile(r'([^\[\]]+)((?:\[\d+\])*)$')
 
 log = logging.getLogger(__name__)
 log.addHandler(NullHandler())
@@ -929,12 +931,16 @@ def batch_get_item(self,
         }
 
         args_map = {}
+        name_placeholders = {}
         if consistent_read:
             args_map[CONSISTENT_READ] = consistent_read
         if return_consumed_capacity:
             operation_kwargs.update(self.get_consumed_capacity_map(return_consumed_capacity))
         if attributes_to_get is not None:
-            args_map[ATTRS_TO_GET] = attributes_to_get
+            projection_expression = self._get_projection_expression(attributes_to_get, name_placeholders)
+            args_map[PROJECTION_EXPRESSION] = projection_expression
+        if name_placeholders:
+            args_map[EXPRESSION_ATTRIBUTE_NAMES] = self._reverse_dict(name_placeholders)
         operation_kwargs[REQUEST_ITEMS][table_name].update(args_map)
 
         keys_map = {KEYS: []}
@@ -958,8 +964,12 @@ def get_item(self,
         Performs the GetItem operation and returns the result
         """
         operation_kwargs = {}
+        name_placeholders = {}
         if attributes_to_get is not None:
-            operation_kwargs[ATTRS_TO_GET] = attributes_to_get
+            projection_expression = self._get_projection_expression(attributes_to_get, name_placeholders)
+            operation_kwargs[PROJECTION_EXPRESSION] = projection_expression
+        if name_placeholders:
+            operation_kwargs[EXPRESSION_ATTRIBUTE_NAMES] = self._reverse_dict(name_placeholders)
         operation_kwargs[CONSISTENT_READ] = consistent_read
         operation_kwargs[TABLE_NAME] = table_name
         operation_kwargs.update(self.get_identifier_map(table_name, hash_key, range_key))
@@ -1129,8 +1139,12 @@ def scan(self,
         Performs the scan operation
         """
         operation_kwargs = {TABLE_NAME: table_name}
+        name_placeholders = {}
         if attributes_to_get is not None:
-            operation_kwargs[ATTRS_TO_GET] = attributes_to_get
+            projection_expression = self._get_projection_expression(attributes_to_get, name_placeholders)
+            operation_kwargs[PROJECTION_EXPRESSION] = projection_expression
+        if name_placeholders:
+            operation_kwargs[EXPRESSION_ATTRIBUTE_NAMES] = self._reverse_dict(name_placeholders)
         if limit is not None:
             operation_kwargs[LIMIT] = limit
         if return_consumed_capacity:
@@ -1183,8 +1197,12 @@ def query(self,
         Performs the Query operation and returns the result
         """
         operation_kwargs = {TABLE_NAME: table_name}
+        name_placeholders = {}
         if attributes_to_get:
-            operation_kwargs[ATTRS_TO_GET] = attributes_to_get
+            projection_expression = self._get_projection_expression(attributes_to_get, name_placeholders)
+            operation_kwargs[PROJECTION_EXPRESSION] = projection_expression
+        if name_placeholders:
+            operation_kwargs[EXPRESSION_ATTRIBUTE_NAMES] = self._reverse_dict(name_placeholders)
         if consistent_read:
             operation_kwargs[CONSISTENT_READ] = True
         if exclusive_start_key:
@@ -1244,6 +1262,15 @@ def query(self,
         except BOTOCORE_EXCEPTIONS as e:
             raise QueryError("Failed to query items: {0}".format(e), e)
 
+    @staticmethod
+    def _get_projection_expression(attributes_to_get, placeholders):
+        expressions = [_substitute_names(attribute, placeholders) for attribute in attributes_to_get]
+        return ', '.join(expressions)
+
+    @staticmethod
+    def _reverse_dict(d):
+        return dict((v, k) for k, v in six.iteritems(d))
+
 
 def _convert_binary(attr):
     if BINARY_SHORT in attr:
@@ -1252,3 +1279,23 @@ def _convert_binary(attr):
         value = attr[BINARY_SET_SHORT]
         if value and len(value):
             attr[BINARY_SET_SHORT] = set(b64decode(v.encode(DEFAULT_ENCODING)) for v in value)
+
+
+def _substitute_names(expression, placeholders):
+    """
+    Replaces names in the given expression with placeholders.
+    Stores the placeholders in the given dictionary.
+    """
+    path_segments = expression.split('.')
+    for idx, segment in enumerate(path_segments):
+        match = PATH_SEGMENT_REGEX.match(segment)
+        if not match:
+            raise ValueError('{0} is not a valid document path'.format(expression))
+        name, indexes = match.groups()
+        if name in placeholders:
+            placeholder = placeholders[name]
+        else:
+            placeholder = '#' + str(len(placeholders))
+            placeholders[name] = placeholder
+        path_segments[idx] = placeholder + indexes
+    return '.'.join(path_segments)
diff --git a/pynamodb/constants.py b/pynamodb/constants.py
@@ -33,7 +33,6 @@
 DELETE_REQUEST = 'DeleteRequest'
 RETURN_VALUES = 'ReturnValues'
 REQUEST_ITEMS = 'RequestItems'
-ATTRS_TO_GET = 'AttributesToGet'
 ATTR_UPDATES = 'AttributeUpdates'
 TABLE_STATUS = 'TableStatus'
 SCAN_FILTER = 'ScanFilter'
@@ -62,6 +61,10 @@
 UTC = 'UTC'
 KEY = 'Key'
 
+# Expression Parameters
+EXPRESSION_ATTRIBUTE_NAMES = 'ExpressionAttributeNames'
+PROJECTION_EXPRESSION = 'ProjectionExpression'
+
 # Defaults
 DEFAULT_ENCODING = 'utf-8'
 DEFAULT_REGION = 'us-east-1'
diff --git a/pynamodb/tests/test_base_connection.py b/pynamodb/tests/test_base_connection.py
@@ -606,7 +606,10 @@ def test_get_item(self):
             )
             params = {
                 'ReturnConsumedCapacity': 'TOTAL',
-                'AttributesToGet': ['ForumName'],
+                'ProjectionExpression': '#0',
+                'ExpressionAttributeNames': {
+                    '#0': 'ForumName'
+                },
                 'Key': {
                     'ForumName': {
                         'S': 'Amazon DynamoDB'
@@ -1231,7 +1234,10 @@ def test_batch_get_item(self):
                 'RequestItems': {
                     'Thread': {
                         'ConsistentRead': True,
-                        'AttributesToGet': ['ForumName'],
+                        'ProjectionExpression': '#0',
+                        'ExpressionAttributeNames': {
+                            '#0': 'ForumName'
+                        },
                         'Keys': [
                             {'ForumName': {'S': 'FooForum'}, 'Subject': {'S': 'thread-0'}},
                             {'ForumName': {'S': 'FooForum'}, 'Subject': {'S': 'thread-1'}},
@@ -1394,7 +1400,10 @@ def test_query(self):
                     }
                 },
                 'IndexName': 'LastPostIndex',
-                'AttributesToGet': ['ForumName'],
+                'ProjectionExpression': '#0',
+                'ExpressionAttributeNames': {
+                    '#0': 'ForumName'
+                },
                 'KeyConditions': {
                     'ForumName': {
                         'ComparisonOperator': 'EQ', 'AttributeValueList': [{
@@ -1650,7 +1659,10 @@ def test_scan(self):
                 attributes_to_get=['ForumName']
             )
             params = {
-                'AttributesToGet': ['ForumName'],
+                'ProjectionExpression': '#0',
+                'ExpressionAttributeNames': {
+                    '#0': 'ForumName'
+                },
                 'ExclusiveStartKey': {
                     "ForumName": {
                         "S": "FooForum"
@@ -2186,3 +2198,25 @@ def test_handle_binary_attributes_for_unprocessed_items(self):
             _assert=True
         )
 
+
+    def test_get_projection_expression(self):
+        attributes_to_get = ['Description', 'RelatedItems[0]', 'ProductReviews.FiveStar']
+        placeholders = {}
+        projection_expression = Connection._get_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0, #1[0], #2.#3"
+        assert placeholders == {'Description': '#0', 'RelatedItems': '#1', 'ProductReviews': '#2', 'FiveStar': '#3'}
+
+
+    def test_get_projection_expression_repeated_names(self):
+        attributes_to_get = ['ProductReviews.FiveStar', 'ProductReviews.ThreeStar', 'ProductReviews.OneStar']
+        placeholders = {}
+        projection_expression = Connection._get_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0.#1, #0.#2, #0.#3"
+        assert placeholders == {'ProductReviews': '#0', 'FiveStar': '#1', 'ThreeStar': '#2', 'OneStar': '#3'}
+
+
+    def test_get_projection_expression_invalid_attribute_raises(self):
+        invalid_attributes = ['', '[0]', 'foo[bar]', 'MyList[-1]', 'MyList[0.4]']
+        for attribute in invalid_attributes:
+            with self.assertRaises(ValueError):
+                Connection._get_projection_expression([attribute], {})
diff --git a/pynamodb/tests/test_model.py b/pynamodb/tests/test_model.py
@@ -2381,7 +2381,10 @@ def test_batch_get(self):
                             {'user_name': {'S': 'hash-1'}},
                             {'user_name': {'S': 'hash-0'}}
                         ],
-                        'AttributesToGet': ['numbers']
+                        'ProjectionExpression': '#0',
+                        'ExpressionAttributeNames': {
+                            '#0': 'numbers'
+                        }
                     }
                 }
             }

Original file line number	Diff line number	Diff line change
`@@ -2381,7 +2381,10 @@ def test_batch_get(self):`
`2381`	`2381`	`{'user_name': {'S': 'hash-1'}},`
`2382`	`2382`	`{'user_name': {'S': 'hash-0'}}`
`2383`	`2383`	`],`
`2384`		`- 'AttributesToGet': ['numbers']`
	`2384`	`+ 'ProjectionExpression': '#0',`
	`2385`	`+ 'ExpressionAttributeNames': {`
	`2386`	`+ '#0': 'numbers'`
	`2387`	`+ }`
`2385`	`2388`	`}`
`2386`	`2389`	`}`
`2387`	`2390`	`}`