Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #2528 from bblommers/feature/2527

Browse Source 
DynamoDB - Allow nested attributes in ProjectionExpressions
This commit is contained in:
Steve Pulec 2019-11-04 22:45:12 -06:00 committed by GitHub
commit f679963ffa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 540 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
moto/dynamodb2/models.py
View File

@ -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):
if attr not in expressions:
result.attrs.pop(attr)
result.filter(projection_expression)
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):
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 +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):
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

33
moto/dynamodb2/responses.py
View File

@ -571,25 +571,22 @@ 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 = ""
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 expression in expression_attribute_names:
projection_expr = (
projection_expr + expression_attribute_names[expression]
)
else:
projection_expr = projection_expr + expression
return projection_expr
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

481
tests/test_dynamodb2/test_dynamodb.py
View File

@ -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")