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

Add ProjectionExpression & ExpressionAttributeNames o DynamoDB get_item & batch_get_item

Browse Source
This commit is contained in:
gruebel 2019-10-08 21:29:09 +02:00
parent e71c06738c
commit cb43796daf
3 changed files with 266 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
moto/dynamodb2/models.py
View File

@ -619,18 +619,29 @@ class Table(BaseModel):
def has_range_key(self):
return self.range_key_attr is not None
def get_item(self, hash_key, range_key=None):
def get_item(self, hash_key, range_key=None, projection_expression=None):
if self.has_range_key and not range_key:
raise ValueError(
"Table has a range key, but no range key was passed into get_item")
try:
result = None
if range_key:
return self.items[hash_key][range_key]
result = self.items[hash_key][range_key]
elif hash_key in self.items:
result = self.items[hash_key]
if hash_key in self.items:
return 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)
raise KeyError
if not result:
raise KeyError
return result
except KeyError:
return None
@ -996,12 +1007,12 @@ class DynamoDBBackend(BaseBackend):
def get_table(self, table_name):
return self.tables.get(table_name)
def get_item(self, table_name, keys):
def get_item(self, table_name, keys, projection_expression=None):
table = self.get_table(table_name)
if not table:
raise ValueError("No table found")
hash_key, range_key = self.get_keys_value(table, keys)
return table.get_item(hash_key, range_key)
return table.get_item(hash_key, range_key, projection_expression)
def query(self, table_name, hash_key_dict, range_comparison, range_value_dicts,
limit, exclusive_start_key, scan_index_forward, projection_expression, index_name=None,

40
moto/dynamodb2/responses.py
View File

@ -305,8 +305,26 @@ class DynamoHandler(BaseResponse):
def get_item(self):
name = self.body['TableName']
key = self.body['Key']
projection_expression = self.body.get('ProjectionExpression')
expression_attribute_names = self.body.get('ExpressionAttributeNames', {})
if projection_expression and expression_attribute_names:
expressions = [x.strip() for x in projection_expression.split(',')]
projection_expression = None
for expression in expressions:
if projection_expression is not None:
projection_expression = projection_expression + ", "
else:
projection_expression = ""
if expression in expression_attribute_names:
projection_expression = projection_expression + \
expression_attribute_names[expression]
else:
projection_expression = projection_expression + expression
try:
item = self.dynamodb_backend.get_item(name, key)
item = self.dynamodb_backend.get_item(name, key, projection_expression)
except ValueError:
er = 'com.amazon.coral.validate#ValidationException'
return self.error(er, 'Validation Exception')
@ -338,9 +356,27 @@ class DynamoHandler(BaseResponse):
er = 'com.amazon.coral.validate#ValidationException'
return self.error(er, 'Provided list of item keys contains duplicates')
attributes_to_get = table_request.get('AttributesToGet')
projection_expression = table_request.get('ProjectionExpression')
expression_attribute_names = table_request.get('ExpressionAttributeNames', {})
if projection_expression and expression_attribute_names:
expressions = [x.strip() for x in projection_expression.split(',')]
projection_expression = None
for expression in expressions:
if projection_expression is not None:
projection_expression = projection_expression + ", "
else:
projection_expression = ""
if expression in expression_attribute_names:
projection_expression = projection_expression + \
expression_attribute_names[expression]
else:
projection_expression = projection_expression + expression
results["Responses"][table_name] = []
for key in keys:
item = self.dynamodb_backend.get_item(table_name, key)
item = self.dynamodb_backend.get_item(table_name, key, projection_expression)
if item:
item_describe = item.describe_attrs(attributes_to_get)
results["Responses"][table_name].append(

212
tests/test_dynamodb2/test_dynamodb.py
View File

@ -369,7 +369,80 @@ def test_query_returns_consumed_capacity():
@mock_dynamodb2
def test_basic_projection_expressions():
def test_basic_projection_expression_using_get_item():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
# Create the DynamoDB table.
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'
})
result = table.get_item(
Key = {
'forum_name': 'the-key',
'subject': '123'
},
ProjectionExpression='body, subject'
)
assert result['Item'] == {
'subject': '123',
'body': 'some test message'
}
# The projection expression should not remove data from storage
result = table.get_item(
Key = {
'forum_name': 'the-key',
'subject': '123'
}
)
assert result['Item'] == {
'forum_name': 'the-key',
'subject': '123',
'body': 'some test message'
}
@mock_dynamodb2
def test_basic_projection_expressions_using_query():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
# Create the DynamoDB table.
@ -452,6 +525,7 @@ def test_basic_projection_expressions():
assert 'body' in results['Items'][1]
assert 'forum_name' in results['Items'][1]
@mock_dynamodb2
def test_basic_projection_expressions_using_scan():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
@ -538,7 +612,73 @@ def test_basic_projection_expressions_using_scan():
@mock_dynamodb2
def test_basic_projection_expressions_with_attr_expression_names():
def test_basic_projection_expression_using_get_item_with_attr_expression_names():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
# Create the DynamoDB table.
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',
'attachment': 'something'
})
table.put_item(Item={
'forum_name': 'not-the-key',
'subject': '123',
'body': 'some other test message',
'attachment': 'something'
})
result = table.get_item(
Key={
'forum_name': 'the-key',
'subject': '123'
},
ProjectionExpression='#rl, #rt, subject',
ExpressionAttributeNames={
'#rl': 'body',
'#rt': 'attachment'
},
)
assert result['Item'] == {
'subject': '123',
'body': 'some test message',
'attachment': 'something'
}
@mock_dynamodb2
def test_basic_projection_expressions_using_query_with_attr_expression_names():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
# Create the DynamoDB table.
@ -603,6 +743,7 @@ def test_basic_projection_expressions_with_attr_expression_names():
assert 'attachment' in results['Items'][0]
assert results['Items'][0]['attachment'] == 'something'
@mock_dynamodb2
def test_basic_projection_expressions_using_scan_with_attr_expression_names():
dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
@ -2233,6 +2374,73 @@ def test_batch_items_returns_all():
assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
@mock_dynamodb2
def test_batch_items_with_basic_projection_expression():
dynamodb = _create_user_table()
returned_items = dynamodb.batch_get_item(RequestItems={
'users': {
'Keys': [{
'username': {'S': 'user0'}
}, {
'username': {'S': 'user1'}
}, {
'username': {'S': 'user2'}
}, {
'username': {'S': 'user3'}
}],
'ConsistentRead': True,
'ProjectionExpression': 'username'
}
})['Responses']['users']
assert len(returned_items) == 3
assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
assert [item.get('foo') for item in returned_items] == [None, None, None]
# The projection expression should not remove data from storage
returned_items = dynamodb.batch_get_item(RequestItems = {
'users': {
'Keys': [{
'username': {'S': 'user0'}
}, {
'username': {'S': 'user1'}
}, {
'username': {'S': 'user2'}
}, {
'username': {'S': 'user3'}
}],
'ConsistentRead': True
}
})['Responses']['users']
assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
assert [item['foo']['S'] for item in returned_items] == ['bar', 'bar', 'bar']
@mock_dynamodb2
def test_batch_items_with_basic_projection_expression_and_attr_expression_names():
dynamodb = _create_user_table()
returned_items = dynamodb.batch_get_item(RequestItems={
'users': {
'Keys': [{
'username': {'S': 'user0'}
}, {
'username': {'S': 'user1'}
}, {
'username': {'S': 'user2'}
}, {
'username': {'S': 'user3'}
}],
'ConsistentRead': True,
'ProjectionExpression': '#rl',
'ExpressionAttributeNames': {
'#rl': 'username'
},
}
})['Responses']['users']
assert len(returned_items) == 3
assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
assert [item.get('foo') for item in returned_items] == [None, None, None]
@mock_dynamodb2
def test_batch_items_should_throw_exception_for_duplicate_request():
client = _create_user_table()