From cb43796dafbf1d7b92745bc2eacaf79915a859f9 Mon Sep 17 00:00:00 2001 From: gruebel Date: Tue, 8 Oct 2019 21:29:09 +0200 Subject: [PATCH] Add ProjectionExpression & ExpressionAttributeNames o DynamoDB get_item & batch_get_item --- moto/dynamodb2/models.py | 25 ++- moto/dynamodb2/responses.py | 40 ++++- tests/test_dynamodb2/test_dynamodb.py | 212 +++++++++++++++++++++++++- 3 files changed, 266 insertions(+), 11 deletions(-) diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py index c06014488..2e4782a7f 100644 --- a/moto/dynamodb2/models.py +++ b/moto/dynamodb2/models.py @@ -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, diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index d07beefd6..da94c4dc1 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -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( diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index b0952f101..fd6681ab6 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' + ) + + 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()