#2527 - DynamoDB - Allow nested attributes in ProjectionExpressions

moto/dynamodb2/models.py

@@ -107,6 +107,19 @@ class DynamoType(object):
             else:
                 self.value.pop(key)
 
+    def filter(self, projection_expressions):
+        nested_projections = [expr[0:expr.index('.')] for expr in projection_expressions if '.' in expr]
+        if self.is_map():
+            expressions_to_delete = []
+            for attr in self.value:
+                if attr not in projection_expressions and attr not in nested_projections:
+                    expressions_to_delete.append(attr)
+                elif attr in nested_projections:
+                    relevant_expressions = [expr[len(attr + '.'):] for expr in projection_expressions if expr.startswith(attr + '.')]
+                    self.value[attr].filter(relevant_expressions)
+            for expr in expressions_to_delete:
+                self.value.pop(expr)
+
     def __hash__(self):
         return hash((self.type, self.value))
 
@@ -477,6 +490,18 @@ class Item(BaseModel):
                     "%s action not support for update_with_attribute_updates" % action
                 )
 
+    # Filter using projection_expression
+    # Ensure a deep copy is used to filter, otherwise actual data will be removed
+    def filter(self, projection_expression):
+        expressions = [x.strip() for x in projection_expression.split(',')]
+        top_level_expressions = [expr[0:expr.index('.')] for expr in expressions if '.' in expr]
+        for attr in list(self.attrs):
+            if attr not in expressions and attr not in top_level_expressions:
+                self.attrs.pop(attr)
+            if attr in top_level_expressions:
+                relevant_expressions = [expr[len(attr + '.'):] for expr in expressions if expr.startswith(attr + '.')]
+                self.attrs[attr].filter(relevant_expressions)
+
 
 class StreamRecord(BaseModel):
     def __init__(self, table, stream_type, event_name, old, new, seq):
@@ -774,11 +799,8 @@ class Table(BaseModel):
                 result = self.items[hash_key]
 
             if projection_expression and result:
-                expressions = [x.strip() for x in projection_expression.split(",")]
                 result = copy.deepcopy(result)
-                for attr in list(result.attrs):
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
+                result.filter(projection_expression)
 
             if not result:
                 raise KeyError
@@ -911,13 +933,10 @@ class Table(BaseModel):
         if filter_expression is not None:
             results = [item for item in results if filter_expression.expr(item)]
 
-        if projection_expression:
-            expressions = [x.strip() for x in projection_expression.split(",")]
         results = copy.deepcopy(results)
+        if projection_expression:
             for result in results:
-                for attr in list(result.attrs):
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
+                result.filter(projection_expression)
 
         results, last_evaluated_key = self._trim_results(
             results, limit, exclusive_start_key
@@ -1004,12 +1023,9 @@ class Table(BaseModel):
                 results.append(item)
 
         if projection_expression:
-            expressions = [x.strip() for x in projection_expression.split(",")]
             results = copy.deepcopy(results)
             for result in results:
-                for attr in list(result.attrs):
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
+                result.filter(projection_expression)
 
         results, last_evaluated_key = self._trim_results(
             results, limit, exclusive_start_key, index_name

moto/dynamodb2/responses.py

@@ -346,9 +346,7 @@ class DynamoHandler(BaseResponse):
         projection_expression = self.body.get("ProjectionExpression")
         expression_attribute_names = self.body.get("ExpressionAttributeNames", {})
 
-        projection_expression = self._adjust_projection_expression(
-            projection_expression, expression_attribute_names
-        )
+        projection_expression = self._adjust_projection_expression(projection_expression, expression_attribute_names)
 
         try:
             item = self.dynamodb_backend.get_item(name, key, projection_expression)
@@ -415,9 +413,7 @@ class DynamoHandler(BaseResponse):
         filter_expression = self.body.get("FilterExpression")
         expression_attribute_values = self.body.get("ExpressionAttributeValues", {})
 
-        projection_expression = self._adjust_projection_expression(
-            projection_expression, expression_attribute_names
-        )
+        projection_expression = self._adjust_projection_expression(projection_expression, expression_attribute_names)
 
         filter_kwargs = {}
 
@@ -571,25 +567,12 @@ class DynamoHandler(BaseResponse):
 
         return dynamo_json_dump(result)
 
-    def _adjust_projection_expression(
-        self, projection_expression, expression_attribute_names
-    ):
-        if projection_expression and expression_attribute_names:
-            expressions = [x.strip() for x in projection_expression.split(",")]
-            projection_expr = None
-            for expression in expressions:
-                if projection_expr is not None:
-                    projection_expr = projection_expr + ", "
-                else:
-                    projection_expr = ""
-
-                if expression in expression_attribute_names:
-                    projection_expr = (
-                        projection_expr + expression_attribute_names[expression]
-                    )
-                else:
-                    projection_expr = projection_expr + expression
-            return projection_expr
+    def _adjust_projection_expression(self, projection_expression, expr_attr_names):
+        def _adjust(expression):
+            return expr_attr_names[expression] if expression in expr_attr_names else expression
+        if projection_expression and expr_attr_names:
+            expressions = [x.strip() for x in projection_expression.split(',')]
+            return ','.join(['.'.join([_adjust(expr) for expr in nested_expr.split('.')]) for nested_expr in expressions])
 
         return projection_expression
 

tests/test_dynamodb2/test_dynamodb.py

@@ -559,6 +559,178 @@ def test_basic_projection_expressions_using_scan():
     assert "forum_name" in results["Items"][1]
 
 
+@mock_dynamodb2
+def test_nested_projection_expression_using_get_item():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a get_item returning all items
+    result = table.get_item(Key={'forum_name': 'key1'},
+                            ProjectionExpression='nested.level1.id, nested.level2')['Item']
+    result.should.equal({'nested': {'level1': {'id': 'id1'},
+                                    'level2': {'id': 'id2', 'include': 'all'}}})
+    # Assert actual data has not been deleted
+    result = table.get_item(Key={'forum_name': 'key1'})['Item']
+    result.should.equal({u'foo': u'bar',
+                         u'forum_name': u'key1',
+                         u'nested': {u'level1': {u'id': u'id1', u'att': u'irrelevant'},
+                                     u'level2': {u'id': u'id2', u'include': u'all'},
+                                     u'level3': {u'id': u'irrelevant'}}})
+
+
+@mock_dynamodb2
+def test_basic_projection_expressions_using_query():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'},
+                                     {'AttributeName': 'subject', 'KeyType': 'RANGE'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'},
+                                                {'AttributeName': 'subject', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'the-key', 'subject': '123', 'body': 'some test message'})
+    table.put_item(Item={'forum_name': 'not-the-key', 'subject': '123', 'body': 'some other test message'})
+
+    # Test a query returning all items
+    result = table.query(KeyConditionExpression=Key('forum_name').eq('the-key'),
+                         ProjectionExpression='body, subject')['Items'][0]
+
+    assert 'body' in result
+    assert result['body'] == 'some test message'
+    assert 'subject' in result
+    assert 'forum_name' not in result
+
+    table.put_item(Item={'forum_name': 'the-key', 'subject': '1234', 'body': 'yet another test message'})
+
+    items = table.query(KeyConditionExpression=Key('forum_name').eq('the-key'),
+                        ProjectionExpression='body')['Items']
+
+    assert 'body' in items[0]
+    assert 'subject' not in items[0]
+    assert items[0]['body'] == 'some test message'
+    assert 'body' in items[1]
+    assert 'subject' not in items[1]
+    assert items[1]['body'] == 'yet another test message'
+
+    # The projection expression should not remove data from storage
+    items = table.query(KeyConditionExpression=Key('forum_name').eq('the-key'))['Items']
+    assert 'subject' in items[0]
+    assert 'body' in items[1]
+    assert 'forum_name' in items[1]
+
+
+@mock_dynamodb2
+def test_nested_projection_expression_using_query():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a query returning all items
+    result = table.query(KeyConditionExpression=Key('forum_name').eq('key1'),
+                         ProjectionExpression="nested.level1.id, nested.level2")['Items'][0]
+
+    assert 'nested' in result
+    result['nested'].should.equal({'level1': {'id': 'id1'},
+                                   'level2': {'id': 'id2', 'include': 'all'}})
+    assert 'foo' not in result
+    # Assert actual data has not been deleted
+    result = table.query(KeyConditionExpression=Key('forum_name').eq('key1'))['Items'][0]
+    result.should.equal({u'foo': u'bar',
+                         u'forum_name': u'key1',
+                         u'nested': {u'level1': {u'id': u'id1', u'att': u'irrelevant'},
+                                     u'level2': {u'id': u'id2', u'include': u'all'},
+                                     u'level3': {u'id': u'irrelevant'}}})
+
+
+@mock_dynamodb2
+def test_basic_projection_expressions_using_scan():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'},
+                                     {'AttributeName': 'subject', 'KeyType': 'RANGE'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'},
+                                                {'AttributeName': 'subject', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+
+    table.put_item(Item={'forum_name': 'the-key', 'subject': '123', 'body': 'some test message'})
+    table.put_item(Item={'forum_name': 'not-the-key', 'subject': '123', 'body': 'some other test message'})
+    # Test a scan returning all items
+    results = table.scan(FilterExpression=Key('forum_name').eq('the-key'),
+                         ProjectionExpression='body, subject')['Items']
+
+    results.should.equal([{u'body': u'some test message', u'subject': u'123'}])
+
+    table.put_item(Item={'forum_name': 'the-key', 'subject': '1234', 'body': 'yet another test message'})
+
+    results = table.scan(FilterExpression=Key('forum_name').eq('the-key'),
+                         ProjectionExpression='body')['Items']
+
+    results.should.equal([{u'body': u'yet another test message'},
+                          {u'body': u'some test message'}])
+
+    # The projection expression should not remove data from storage
+    results = table.query(KeyConditionExpression=Key('forum_name').eq('the-key'))
+    assert 'subject' in results['Items'][0]
+    assert 'body' in results['Items'][1]
+    assert 'forum_name' in results['Items'][1]
+
+
+@mock_dynamodb2
+def test_nested_projection_expression_using_scan():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a scan
+    results = table.scan(FilterExpression=Key('forum_name').eq('key1'),
+                         ProjectionExpression='nested.level1.id, nested.level2')['Items']
+    results.should.equal([{'nested': {'level1': {'id': 'id1'},
+                                      'level2': {'include': 'all', 'id': 'id2'}}}])
+    # Assert original data is still there
+    results = table.scan(FilterExpression=Key('forum_name').eq('key1'))['Items']
+    results.should.equal([{'forum_name': 'key1',
+                           'foo': 'bar',
+                           'nested': {'level1': {'att': 'irrelevant', 'id': 'id1'},
+                                      'level2': {'include': 'all', 'id': 'id2'},
+                                      'level3': {'id': 'irrelevant'}}}])
+
+
 @mock_dynamodb2
 def test_basic_projection_expression_using_get_item_with_attr_expression_names():
     dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
@@ -658,6 +830,71 @@ def test_basic_projection_expressions_using_query_with_attr_expression_names():
     assert results["Items"][0]["attachment"] == "something"
 
 
+@mock_dynamodb2
+def test_nested_projection_expression_using_get_item_with_attr_expression():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a get_item returning all items
+    result = table.get_item(Key={'forum_name': 'key1'},
+                            ProjectionExpression='#nst.level1.id, #nst.#lvl2',
+                            ExpressionAttributeNames={'#nst': 'nested', '#lvl2': 'level2'})['Item']
+    result.should.equal({'nested': {'level1': {'id': 'id1'},
+                                    'level2': {'id': 'id2', 'include': 'all'}}})
+    # Assert actual data has not been deleted
+    result = table.get_item(Key={'forum_name': 'key1'})['Item']
+    result.should.equal({u'foo': u'bar',
+                         u'forum_name': u'key1',
+                         u'nested': {u'level1': {u'id': u'id1', u'att': u'irrelevant'},
+                                     u'level2': {u'id': u'id2', u'include': u'all'},
+                                     u'level3': {u'id': u'irrelevant'}}})
+
+
+@mock_dynamodb2
+def test_nested_projection_expression_using_query_with_attr_expression_names():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a query returning all items
+    result = table.query(KeyConditionExpression=Key('forum_name').eq('key1'),
+                         ProjectionExpression="#nst.level1.id, #nst.#lvl2",
+                         ExpressionAttributeNames={'#nst': 'nested', '#lvl2': 'level2'})['Items'][0]
+
+    assert 'nested' in result
+    result['nested'].should.equal({'level1': {'id': 'id1'},
+                                   'level2': {'id': 'id2', 'include': 'all'}})
+    assert 'foo' not in result
+    # Assert actual data has not been deleted
+    result = table.query(KeyConditionExpression=Key('forum_name').eq('key1'))['Items'][0]
+    result.should.equal( {u'foo': u'bar',
+                          u'forum_name': u'key1',
+                          u'nested': {u'level1': {u'id': u'id1', u'att': u'irrelevant'},
+                                      u'level2': {u'id': u'id2', u'include': u'all'},
+                                      u'level3': {u'id': u'irrelevant'}}})
+
+
 @mock_dynamodb2
 def test_basic_projection_expressions_using_scan_with_attr_expression_names():
     dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
@@ -719,6 +956,37 @@ def test_basic_projection_expressions_using_scan_with_attr_expression_names():
     assert "form_name" not in results["Items"][0]
 
 
+@mock_dynamodb2
+def test_nested_projection_expression_using_scan_with_attr_expression_names():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    dynamodb.create_table(TableName='users',
+                          KeySchema=[{'AttributeName': 'forum_name', 'KeyType': 'HASH'}],
+                          AttributeDefinitions=[{'AttributeName': 'forum_name', 'AttributeType': 'S'}],
+                          ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5})
+    table = dynamodb.Table('users')
+    table.put_item(Item={'forum_name': 'key1', 'nested': {'level1': {'id': 'id1', 'att': 'irrelevant'},
+                                                          'level2': {'id': 'id2', 'include': 'all'},
+                                                          'level3': {'id': 'irrelevant'}},
+                         'foo': 'bar'})
+    table.put_item(Item={'forum_name': 'key2', 'nested': {'id': 'id2', 'incode': 'code2'}, 'foo': 'bar'})
+
+    # Test a scan
+    results = table.scan(FilterExpression=Key('forum_name').eq('key1'),
+                         ProjectionExpression='nested.level1.id, nested.level2',
+                         ExpressionAttributeNames={'#nst': 'nested', '#lvl2': 'level2'})['Items']
+    results.should.equal([{'nested': {'level1': {'id': 'id1'},
+                                      'level2': {'include': 'all', 'id': 'id2'}}}])
+    # Assert original data is still there
+    results = table.scan(FilterExpression=Key('forum_name').eq('key1'))['Items']
+    results.should.equal([{'forum_name': 'key1',
+                           'foo': 'bar',
+                           'nested': {'level1': {'att': 'irrelevant', 'id': 'id1'},
+                                      'level2': {'include': 'all', 'id': 'id2'},
+                                      'level3': {'id': 'irrelevant'}}}])
+
+
 @mock_dynamodb2
 def test_put_item_returns_consumed_capacity():
     dynamodb = boto3.resource("dynamodb", region_name="us-east-1")