Merge pull request #2528 from bblommers/feature/2527

DynamoDB - Allow nested attributes in ProjectionExpressions

moto/dynamodb2/models.py

@@ -107,6 +107,28 @@ 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 +499,24 @@ 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 +814,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):
+                result.filter(projection_expression)
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
 
             if not result:
                 raise KeyError
@@ -911,13 +948,10 @@ class Table(BaseModel):
         if filter_expression is not None:
             results = [item for item in results if filter_expression.expr(item)]
 
+        results = copy.deepcopy(results)
         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):
+                result.filter(projection_expression)
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
 
         results, last_evaluated_key = self._trim_results(
             results, limit, exclusive_start_key
@@ -1004,12 +1038,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):
+                result.filter(projection_expression)
-                    if attr not in expressions:
-                        result.attrs.pop(attr)
 
         results, last_evaluated_key = self._trim_results(
             results, limit, exclusive_start_key, index_name

moto/dynamodb2/responses.py

@@ -571,25 +571,22 @@ class DynamoHandler(BaseResponse):
 
         return dynamo_json_dump(result)
 
-    def _adjust_projection_expression(
+    def _adjust_projection_expression(self, projection_expression, expr_attr_names):
-        self, projection_expression, expression_attribute_names
+        def _adjust(expression):
-    ):
+            return (
-        if projection_expression and expression_attribute_names:
+                expr_attr_names[expression]
-            expressions = [x.strip() for x in projection_expression.split(",")]
+                if expression in expr_attr_names
-            projection_expr = None
+                else expression
-            for expression in expressions:
+            )
-                if projection_expr is not None:
-                    projection_expr = projection_expr + ", "
-                else:
-                    projection_expr = ""
 
-                if expression in expression_attribute_names:
+        if projection_expression and expr_attr_names:
-                    projection_expr = (
+            expressions = [x.strip() for x in projection_expression.split(",")]
-                        projection_expr + expression_attribute_names[expression]
+            return ",".join(
-                    )
+                [
-                else:
+                    ".".join([_adjust(expr) for expr in nested_expr.split(".")])
-                    projection_expr = projection_expr + expression
+                    for nested_expr in expressions
-            return projection_expr
+                ]
+            )
 
         return projection_expression
 

tests/test_dynamodb2/test_dynamodb.py

@@ -559,6 +559,308 @@ 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(
+        {
+            "foo": "bar",
+            "forum_name": "key1",
+            "nested": {
+                "level1": {"id": "id1", "att": "irrelevant"},
+                "level2": {"id": "id2", "include": "all"},
+                "level3": {"id": "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(
+        {
+            "foo": "bar",
+            "forum_name": "key1",
+            "nested": {
+                "level1": {"id": "id1", "att": "irrelevant"},
+                "level2": {"id": "id2", "include": "all"},
+                "level3": {"id": "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([{"body": "some test message", "subject": "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"]
+
+    assert {"body": "some test message"} in results
+    assert {"body": "yet another test message"} in results
+
+    # 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 +960,121 @@ 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(
+        {
+            "foo": "bar",
+            "forum_name": "key1",
+            "nested": {
+                "level1": {"id": "id1", "att": "irrelevant"},
+                "level2": {"id": "id2", "include": "all"},
+                "level3": {"id": "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(
+        {
+            "foo": "bar",
+            "forum_name": "key1",
+            "nested": {
+                "level1": {"id": "id1", "att": "irrelevant"},
+                "level2": {"id": "id2", "include": "all"},
+                "level3": {"id": "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 +1136,70 @@ 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")