diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py index 5915d6eea..fde269726 100644 --- a/moto/dynamodb2/models.py +++ b/moto/dynamodb2/models.py @@ -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) diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index cf715bfbc..37b73160e 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -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] diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index 8f320cbab..2d58740f5 100644 --- a/tests/test_dynamodb2/test_dynamodb.py +++ b/tests/test_dynamodb2/test_dynamodb.py @@ -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'