Add KeyConditionExpression to dynamo.

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,
+    'EQ': EQ_FUNCTION,
-    'NE': lambda item_value, test_value: item_value != test_value,
+    '=': EQ_FUNCTION,
-    'LE': lambda item_value, test_value: item_value <= test_value,
+
-    'LT': lambda item_value, test_value: item_value < test_value,
+    'NE': NE_FUNCTION,
-    'GE': lambda item_value, test_value: item_value >= test_value,
+    '!=': NE_FUNCTION,
-    'GT': lambda item_value, test_value: item_value > test_value,
+
+    '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,

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())
+        # {u'KeyConditionExpression': u'#n0 = :v0', u'ExpressionAttributeValues': {u':v0': {u'S': u'johndoe'}}, u'ExpressionAttributeNames': {u'#n0': u'username'}}
-        # hash_key_name, range_key_name = dynamodb_backend2.get_table_keys_name(name)
+        key_condition_expression = self.body.get('KeyConditionExpression')
-        if hash_key_name is None:
+        if key_condition_expression:
-            er = "'com.amazonaws.dynamodb.v20120810#ResourceNotFoundException"
+            value_alias_map = self.body['ExpressionAttributeValues']
-            return self.error(er)
+
-        hash_key = key_conditions[hash_key_name]['AttributeValueList'][0]
+            if " AND " in key_condition_expression:
-        if len(key_conditions) == 1:
+                expressions = key_condition_expression.split(" AND ", 1)
-            range_comparison = None
+                hash_key_expression = expressions[0]
-            range_values = []
+                # TODO implement more than one range expression and OR operators
-        else:
+                range_key_expression = expressions[1].replace(")", "")
-            if range_key_name is None:
+                range_key_expression_components = range_key_expression.split()
-                er = "com.amazon.coral.validate#ValidationException"
+                range_comparison = range_key_expression_components[1]
-                return self.error(er)
+                if 'AND' in range_key_expression:
-            else:
+                    range_comparison = 'BETWEEN'
-                range_condition = key_conditions[range_key_name]
+                    range_values = [
-                if range_condition:
+                        value_alias_map[range_key_expression_components[2]],
-                    range_comparison = range_condition['ComparisonOperator']
+                        value_alias_map[range_key_expression_components[4]],
-                    range_values = range_condition['AttributeValueList']
+                    ]
+                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 = {

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'}
+    data2 = {'id': '123', 'range': 'abc', 'data': '345'}
-    table.put_item(data=data2, overwrite = True)
+    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)

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'}
+    data2 = {'id': '123', 'data': '345'}
-    table.put_item(data=data2, overwrite = True)
+    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)
+    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"})