From e4408152d196eafc8245e79b771d54b094ec212d Mon Sep 17 00:00:00 2001 From: Steve Pulec Date: Sat, 1 Aug 2015 19:32:33 -0400 Subject: [PATCH] Add KeyConditionExpression to dynamo. --- moto/dynamodb2/comparisons.py | 32 ++++- moto/dynamodb2/responses.py | 76 ++++++++--- .../test_dynamodb_table_with_range_key.py | 128 ++++++++++++++++-- .../test_dynamodb_table_without_range_key.py | 67 ++++++--- 4 files changed, 246 insertions(+), 57 deletions(-) diff --git a/moto/dynamodb2/comparisons.py b/moto/dynamodb2/comparisons.py index 86f582179..808e120bc 100644 --- a/moto/dynamodb2/comparisons.py +++ b/moto/dynamodb2/comparisons.py @@ -1,12 +1,32 @@ from __future__ import unicode_literals # TODO add tests for all of these + +EQ_FUNCTION = lambda item_value, test_value: item_value == test_value +NE_FUNCTION = lambda item_value, test_value: item_value != test_value +LE_FUNCTION = lambda item_value, test_value: item_value <= test_value +LT_FUNCTION = lambda item_value, test_value: item_value < test_value +GE_FUNCTION = lambda item_value, test_value: item_value >= test_value +GT_FUNCTION = lambda item_value, test_value: item_value > test_value + COMPARISON_FUNCS = { - 'EQ': lambda item_value, test_value: item_value == test_value, - 'NE': lambda item_value, test_value: item_value != test_value, - 'LE': lambda item_value, test_value: item_value <= test_value, - 'LT': lambda item_value, test_value: item_value < test_value, - 'GE': lambda item_value, test_value: item_value >= test_value, - 'GT': lambda item_value, test_value: item_value > test_value, + 'EQ': EQ_FUNCTION, + '=': EQ_FUNCTION, + + 'NE': NE_FUNCTION, + '!=': NE_FUNCTION, + + 'LE': LE_FUNCTION, + '<=': LE_FUNCTION, + + 'LT': LT_FUNCTION, + '<': LT_FUNCTION, + + 'GE': GE_FUNCTION, + '>=': GE_FUNCTION, + + 'GT': GT_FUNCTION, + '>': GT_FUNCTION, + 'NULL': lambda item_value: item_value is None, 'NOT_NULL': lambda item_value: item_value is not None, 'CONTAINS': lambda item_value, test_value: test_value in item_value, diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index 81b6a7f47..069d31229 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -228,28 +228,64 @@ class DynamoHandler(BaseResponse): def query(self): name = self.body['TableName'] - key_conditions = self.body['KeyConditions'] - hash_key_name, range_key_name = dynamodb_backend2.get_table_keys_name(name, key_conditions.keys()) - # hash_key_name, range_key_name = dynamodb_backend2.get_table_keys_name(name) - if hash_key_name is None: - er = "'com.amazonaws.dynamodb.v20120810#ResourceNotFoundException" - return self.error(er) - hash_key = key_conditions[hash_key_name]['AttributeValueList'][0] - if len(key_conditions) == 1: - range_comparison = None - range_values = [] - else: - if range_key_name is None: - er = "com.amazon.coral.validate#ValidationException" - return self.error(er) - else: - range_condition = key_conditions[range_key_name] - if range_condition: - range_comparison = range_condition['ComparisonOperator'] - range_values = range_condition['AttributeValueList'] + + # {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') + if key_condition_expression: + value_alias_map = self.body['ExpressionAttributeValues'] + + if " AND " in key_condition_expression: + expressions = key_condition_expression.split(" AND ", 1) + hash_key_expression = expressions[0] + # TODO implement more than one range expression and OR operators + range_key_expression = expressions[1].replace(")", "") + range_key_expression_components = range_key_expression.split() + range_comparison = range_key_expression_components[1] + if 'AND' in range_key_expression: + range_comparison = 'BETWEEN' + range_values = [ + value_alias_map[range_key_expression_components[2]], + value_alias_map[range_key_expression_components[4]], + ] + elif 'begins_with' in range_key_expression: + range_comparison = 'BEGINS_WITH' + range_values = [ + value_alias_map[range_key_expression_components[1]], + ] else: + range_values = [value_alias_map[range_key_expression_components[2]]] + else: + hash_key_expression = key_condition_expression + range_comparison = None + range_values = [] + + hash_key_value_alias = hash_key_expression.split("=")[1].strip() + hash_key = value_alias_map[hash_key_value_alias] + else: + # 'KeyConditions': {u'forum_name': {u'ComparisonOperator': u'EQ', u'AttributeValueList': [{u'S': u'the-key'}]}} + key_conditions = self.body.get('KeyConditions') + if key_conditions: + hash_key_name, range_key_name = dynamodb_backend2.get_table_keys_name(name, key_conditions.keys()) + if hash_key_name is None: + er = "'com.amazonaws.dynamodb.v20120810#ResourceNotFoundException" + return self.error(er) + hash_key = key_conditions[hash_key_name]['AttributeValueList'][0] + if len(key_conditions) == 1: range_comparison = None range_values = [] + else: + if range_key_name is None: + er = "com.amazon.coral.validate#ValidationException" + return self.error(er) + else: + range_condition = key_conditions[range_key_name] + if range_condition: + range_comparison = range_condition['ComparisonOperator'] + range_values = range_condition['AttributeValueList'] + else: + range_comparison = None + range_values = [] + items, last_page = dynamodb_backend2.query(name, hash_key, range_comparison, range_values) if items is None: er = 'com.amazonaws.dynamodb.v20111205#ResourceNotFoundException' @@ -260,7 +296,7 @@ class DynamoHandler(BaseResponse): items = items[:limit] reversed = self.body.get("ScanIndexForward") - if reversed is not False: + if reversed is False: items.reverse() result = { diff --git a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py index e7ae60e7e..12e3aa15b 100644 --- a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py +++ b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py @@ -1,6 +1,8 @@ from __future__ import unicode_literals import boto +import boto3 +from boto3.dynamodb.conditions import Key import sure # noqa from freezegun import freeze_time from moto import mock_dynamodb2 @@ -253,31 +255,31 @@ def test_query(): table.count().should.equal(4) - results = table.query(forum_name__eq='the-key', subject__gt='1', consistent=True) + results = table.query_2(forum_name__eq='the-key', subject__gt='1', consistent=True) expected = ["123", "456", "789"] for index, item in enumerate(results): item["subject"].should.equal(expected[index]) - results = table.query(forum_name__eq="the-key", subject__gt='1', reverse=True) + results = table.query_2(forum_name__eq="the-key", subject__gt='1', reverse=True) for index, item in enumerate(results): item["subject"].should.equal(expected[len(expected) - 1 - index]) - results = table.query(forum_name__eq='the-key', subject__gt='1', consistent=True) + results = table.query_2(forum_name__eq='the-key', subject__gt='1', consistent=True) sum(1 for _ in results).should.equal(3) - results = table.query(forum_name__eq='the-key', subject__gt='234', consistent=True) + results = table.query_2(forum_name__eq='the-key', subject__gt='234', consistent=True) sum(1 for _ in results).should.equal(2) - results = table.query(forum_name__eq='the-key', subject__gt='9999') + results = table.query_2(forum_name__eq='the-key', subject__gt='9999') sum(1 for _ in results).should.equal(0) - results = table.query(forum_name__eq='the-key', subject__beginswith='12') + results = table.query_2(forum_name__eq='the-key', subject__beginswith='12') sum(1 for _ in results).should.equal(1) - results = table.query(forum_name__eq='the-key', subject__beginswith='7') + results = table.query_2(forum_name__eq='the-key', subject__beginswith='7') sum(1 for _ in results).should.equal(1) - results = table.query(forum_name__eq='the-key', subject__between=['567', '890']) + results = table.query_2(forum_name__eq='the-key', subject__between=['567', '890']) sum(1 for _ in results).should.equal(1) @@ -558,7 +560,6 @@ def test_lookup(): @mock_dynamodb2 def test_failed_overwrite(): - from decimal import Decimal table = Table.create('messages', schema=[ HashKey('id'), RangeKey('range'), @@ -567,19 +568,19 @@ def test_failed_overwrite(): 'write': 3, }) - data1 = {'id': '123', 'range': 'abc', 'data':'678'} + data1 = {'id': '123', 'range': 'abc', 'data': '678'} table.put_item(data=data1) - data2 = {'id': '123', 'range': 'abc', 'data':'345'} - table.put_item(data=data2, overwrite = True) + data2 = {'id': '123', 'range': 'abc', 'data': '345'} + table.put_item(data=data2, overwrite=True) - data3 = {'id': '123', 'range': 'abc', 'data':'812'} + data3 = {'id': '123', 'range': 'abc', 'data': '812'} table.put_item.when.called_with(data=data3).should.throw(ConditionalCheckFailedException) returned_item = table.lookup('123', 'abc') dict(returned_item).should.equal(data2) - data4 = {'id': '123', 'range': 'ghi', 'data':812} + data4 = {'id': '123', 'range': 'ghi', 'data': 812} table.put_item(data=data4) returned_item = table.lookup('123', 'ghi') @@ -593,7 +594,7 @@ def test_conflicting_writes(): RangeKey('range'), ]) - item_data = {'id': '123', 'range':'abc', 'data':'678'} + item_data = {'id': '123', 'range': 'abc', 'data': '678'} item1 = Item(table, item_data) item2 = Item(table, item_data) item1.save() @@ -603,3 +604,100 @@ def test_conflicting_writes(): item1.save() item2.save.when.called_with().should.throw(ConditionalCheckFailedException) + +""" +boto3 +""" + + +@mock_dynamodb2 +def test_boto3_conditions(): + 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' + }) + table.put_item(Item={ + 'forum_name': 'the-key', + 'subject': '456' + }) + table.put_item(Item={ + 'forum_name': 'the-key', + 'subject': '789' + }) + + # Test a query returning all items + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").gt('1'), + ScanIndexForward=True, + ) + expected = ["123", "456", "789"] + for index, item in enumerate(results['Items']): + item["subject"].should.equal(expected[index]) + + # Return all items again, but in reverse + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").gt('1'), + ScanIndexForward=False, + ) + for index, item in enumerate(reversed(results['Items'])): + item["subject"].should.equal(expected[index]) + + # Filter the subjects to only return some of the results + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").gt('234'), + ConsistentRead=True, + ) + results['Count'].should.equal(2) + + # Filter to return no results + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").gt('9999') + ) + results['Count'].should.equal(0) + + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").begins_with('12') + ) + results['Count'].should.equal(1) + + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").begins_with('7') + ) + results['Count'].should.equal(1) + + results = table.query( + KeyConditionExpression=Key('forum_name').eq('the-key') & Key("subject").between('567', '890') + ) + results['Count'].should.equal(1) diff --git a/tests/test_dynamodb2/test_dynamodb_table_without_range_key.py b/tests/test_dynamodb2/test_dynamodb_table_without_range_key.py index 1710bd650..808805b8d 100644 --- a/tests/test_dynamodb2/test_dynamodb_table_without_range_key.py +++ b/tests/test_dynamodb2/test_dynamodb_table_without_range_key.py @@ -1,6 +1,8 @@ from __future__ import unicode_literals import boto +import boto3 +from boto3.dynamodb.conditions import Key import sure # noqa from freezegun import freeze_time from boto.exception import JSONResponseError @@ -135,14 +137,6 @@ def test_item_put_without_table(): ).should.throw(JSONResponseError) -@requires_boto_gte("2.9") -@mock_dynamodb2 -def test_get_missing_item(): - table = create_table() - - table.get_item.when.called_with(test_hash=3241526475).should.throw(JSONResponseError) - - @requires_boto_gte("2.9") @mock_dynamodb2 def test_get_item_with_undeclared_table(): @@ -449,7 +443,6 @@ def test_update_item_set(): @mock_dynamodb2 def test_failed_overwrite(): - from decimal import Decimal table = Table.create('messages', schema=[ HashKey('id'), ], throughput={ @@ -457,19 +450,19 @@ def test_failed_overwrite(): 'write': 3, }) - data1 = {'id': '123', 'data':'678'} + data1 = {'id': '123', 'data': '678'} table.put_item(data=data1) - data2 = {'id': '123', 'data':'345'} - table.put_item(data=data2, overwrite = True) + data2 = {'id': '123', 'data': '345'} + table.put_item(data=data2, overwrite=True) - data3 = {'id': '123', 'data':'812'} + data3 = {'id': '123', 'data': '812'} table.put_item.when.called_with(data=data3).should.throw(ConditionalCheckFailedException) returned_item = table.lookup('123') dict(returned_item).should.equal(data2) - data4 = {'id': '124', 'data':812} + data4 = {'id': '124', 'data': 812} table.put_item(data=data4) returned_item = table.lookup('124') @@ -482,7 +475,7 @@ def test_conflicting_writes(): HashKey('id'), ]) - item_data = {'id': '123', 'data':'678'} + item_data = {'id': '123', 'data': '678'} item1 = Item(table, item_data) item2 = Item(table, item_data) item1.save() @@ -491,4 +484,46 @@ def test_conflicting_writes(): item2['data'] = '912' item1.save() - item2.save.when.called_with().should.throw(ConditionalCheckFailedException) \ No newline at end of file + item2.save.when.called_with().should.throw(ConditionalCheckFailedException) + + +""" +boto3 +""" + + +@mock_dynamodb2 +def test_boto3_conditions(): + dynamodb = boto3.resource('dynamodb', region_name='us-east-1') + + # Create the DynamoDB table. + table = dynamodb.create_table( + TableName='users', + KeySchema=[ + { + 'AttributeName': 'username', + 'KeyType': 'HASH' + }, + ], + AttributeDefinitions=[ + { + 'AttributeName': 'username', + 'AttributeType': 'S' + }, + ], + ProvisionedThroughput={ + 'ReadCapacityUnits': 5, + 'WriteCapacityUnits': 5 + } + ) + table = dynamodb.Table('users') + + table.put_item(Item={'username': 'johndoe'}) + table.put_item(Item={'username': 'janedoe'}) + + response = table.query( + KeyConditionExpression=Key('username').eq('johndoe') + ) + response['Count'].should.equal(1) + response['Items'].should.have.length_of(1) + response['Items'][0].should.equal({"username": "johndoe"})