diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py index efd784ab0..630d031f2 100644 --- a/moto/dynamodb2/models.py +++ b/moto/dynamodb2/models.py @@ -682,18 +682,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 @@ -1059,12 +1070,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, diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index e3ce1b39e..683d23c2f 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -308,8 +308,15 @@ 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', {}) + + projection_expression = self._adjust_projection_expression( + projection_expression, expression_attribute_names + ) + 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') @@ -341,9 +348,16 @@ 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', {}) + + projection_expression = self._adjust_projection_expression( + projection_expression, expression_attribute_names + ) + 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( @@ -373,20 +387,9 @@ class DynamoHandler(BaseResponse): filter_expression = self.body.get('FilterExpression') expression_attribute_values = self.body.get('ExpressionAttributeValues', {}) - 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 + projection_expression = self._adjust_projection_expression( + projection_expression, expression_attribute_names + ) filter_kwargs = {} @@ -522,6 +525,25 @@ 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 = "" + + if expression in expression_attribute_names: + projection_expr = projection_expr + \ + expression_attribute_names[expression] + else: + projection_expr = projection_expr + expression + return projection_expr + + return projection_expression + def scan(self): name = self.body['TableName'] diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index 8aa1480d5..749b30c26 100644 --- a/tests/test_dynamodb2/test_dynamodb.py +++ b/tests/test_dynamodb2/test_dynamodb.py @@ -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' + ) + + result['Item'].should.be.equal({ + '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' + } + ) + + result['Item'].should.be.equal({ + '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' + }, + ) + + result['Item'].should.be.equal({ + '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') @@ -2250,6 +2391,76 @@ 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'] + + returned_items.should.have.length_of(3) + [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3']) + [item.get('foo') for item in returned_items].should.be.equal([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'] + + [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3']) + [item['foo']['S'] for item in returned_items].should.be.equal(['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'] + + returned_items.should.have.length_of(3) + [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3']) + [item.get('foo') for item in returned_items].should.be.equal([None, None, None]) + + @mock_dynamodb2 def test_batch_items_should_throw_exception_for_duplicate_request(): client = _create_user_table()