Merge pull request #1176 from dbfr3qs/master

Add basic ProjectionExpression to dynamodb2.query
This commit is contained in:
Jack Danger 2017-09-21 22:28:31 -07:00 committed by GitHub
commit 79ffdd8b7d
3 changed files with 216 additions and 5 deletions

View File

@ -412,7 +412,8 @@ class Table(BaseModel):
return None
def query(self, hash_key, range_comparison, range_objs, limit,
exclusive_start_key, scan_index_forward, index_name=None, **filter_kwargs):
exclusive_start_key, scan_index_forward, projection_expression,
index_name=None, **filter_kwargs):
results = []
if index_name:
all_indexes = (self.global_indexes or []) + (self.indexes or [])
@ -483,6 +484,13 @@ class Table(BaseModel):
else:
results.sort(key=lambda item: item.range_key)
if projection_expression:
expressions = [x.strip() for x in projection_expression.split(',')]
for result in possible_results:
for attr in list(result.attrs):
if attr not in expressions:
result.attrs.pop(attr)
if scan_index_forward is False:
results.reverse()
@ -678,7 +686,7 @@ class DynamoDBBackend(BaseBackend):
return table.get_item(hash_key, range_key)
def query(self, table_name, hash_key_dict, range_comparison, range_value_dicts,
limit, exclusive_start_key, scan_index_forward, index_name=None, **filter_kwargs):
limit, exclusive_start_key, scan_index_forward, projection_expression, index_name=None, **filter_kwargs):
table = self.tables.get(table_name)
if not table:
return None, None
@ -688,7 +696,7 @@ class DynamoDBBackend(BaseBackend):
for range_value in range_value_dicts]
return table.query(hash_key, range_comparison, range_values, limit,
exclusive_start_key, scan_index_forward, index_name, **filter_kwargs)
exclusive_start_key, scan_index_forward, projection_expression, index_name, **filter_kwargs)
def scan(self, table_name, filters, limit, exclusive_start_key):
table = self.tables.get(table_name)

View File

@ -276,6 +276,15 @@ class DynamoHandler(BaseResponse):
name = self.body['TableName']
# {u'KeyConditionExpression': u'#n0 = :v0', u'ExpressionAttributeValues': {u':v0': {u'S': u'johndoe'}}, u'ExpressionAttributeNames': {u'#n0': u'username'}}
key_condition_expression = self.body.get('KeyConditionExpression')
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(',')]
for expression in expressions:
if expression in expression_attribute_names:
projection_expression = projection_expression.replace(expression, expression_attribute_names[expression])
filter_kwargs = {}
if key_condition_expression:
value_alias_map = self.body['ExpressionAttributeValues']
@ -383,16 +392,20 @@ class DynamoHandler(BaseResponse):
scan_index_forward = self.body.get("ScanIndexForward")
items, scanned_count, last_evaluated_key = dynamodb_backend2.query(
name, hash_key, range_comparison, range_values, limit,
exclusive_start_key, scan_index_forward, index_name=index_name, **filter_kwargs)
exclusive_start_key, scan_index_forward, projection_expression, index_name=index_name, **filter_kwargs)
if items is None:
er = 'com.amazonaws.dynamodb.v20111205#ResourceNotFoundException'
return self.error(er, 'Requested resource not found')
result = {
"Count": len(items),
"ConsumedCapacityUnits": 1,
'ConsumedCapacity': {
'TableName': name,
'CapacityUnits': 1,
},
"ScannedCount": scanned_count
}
if self.body.get('Select', '').upper() != 'COUNT':
result["Items"] = [item.attrs for item in items]

View File

@ -9,6 +9,7 @@ from moto import mock_dynamodb2, mock_dynamodb2_deprecated
from moto.dynamodb2 import dynamodb_backend2
from boto.exception import JSONResponseError
from botocore.exceptions import ClientError
from boto3.dynamodb.conditions import Key
from tests.helpers import requires_boto_gte
import tests.backport_assert_raises
from nose.tools import assert_raises
@ -228,3 +229,192 @@ def test_scan_returns_consumed_capacity():
assert 'ConsumedCapacity' in response
assert 'CapacityUnits' in response['ConsumedCapacity']
assert response['ConsumedCapacity']['TableName'] == name
@requires_boto_gte("2.9")
@mock_dynamodb2
def test_query_returns_consumed_capacity():
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'
})
results = table.query(
KeyConditionExpression=Key('forum_name').eq(
'the-key')
)
assert 'ConsumedCapacity' in results
assert 'CapacityUnits' in results['ConsumedCapacity']
assert results['ConsumedCapacity']['CapacityUnits'] == 1
@mock_dynamodb2
def test_basic_projection_expressions():
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'
})
# Test a query returning all items
results = table.query(
KeyConditionExpression=Key('forum_name').eq(
'the-key'),
ProjectionExpression='body, subject'
)
assert 'body' in results['Items'][0]
assert results['Items'][0]['body'] == 'some test message'
assert 'subject' in results['Items'][0]
table.put_item(Item={
'forum_name': 'the-key',
'subject': '1234',
'body': 'yet another test message'
})
results = table.query(
KeyConditionExpression=Key('forum_name').eq(
'the-key'),
ProjectionExpression='body'
)
assert 'body' in results['Items'][0]
assert results['Items'][0]['body'] == 'some test message'
assert 'body' in results['Items'][1]
assert results['Items'][1]['body'] == 'yet another test message'
@mock_dynamodb2
def test_basic_projection_expressions_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'
})
# Test a query returning all items
results = table.query(
KeyConditionExpression=Key('forum_name').eq(
'the-key'),
ProjectionExpression='#rl, #rt, subject',
ExpressionAttributeNames={
'#rl': 'body',
'#rt': 'attachment'
},
)
assert 'body' in results['Items'][0]
assert results['Items'][0]['body'] == 'some test message'
assert 'subject' in results['Items'][0]
assert results['Items'][0]['subject'] == '123'
assert 'attachment' in results['Items'][0]
assert results['Items'][0]['attachment'] == 'something'