Merge pull request #567 from im-auld/query-filters-issue-164

Query filters issue 164

moto/dynamodb2/models.py

@@ -309,7 +309,7 @@ class Table(object):
             return None
 
     def query(self, hash_key, range_comparison, range_objs, limit,
-              exclusive_start_key, scan_index_forward, index_name=None):
+              exclusive_start_key, scan_index_forward, index_name=None, **filter_kwargs):
         results = []
 
         if index_name:
@@ -354,8 +354,16 @@ class Table(object):
                 for result in possible_results:
                     if result.range_key.compare(range_comparison, range_objs):
                         results.append(result)
-        else:
+
-            # If we're not filtering on range key, return all values
+        if filter_kwargs:
+            for result in possible_results:
+                for field, value in filter_kwargs.items():
+                    dynamo_types = [DynamoType(ele) for ele in value["AttributeValueList"]]
+                    if result.attrs.get(field).compare(value['ComparisonOperator'], dynamo_types):
+                        results.append(result)
+
+        if not range_comparison and not filter_kwargs:
+            # If we're not filtering on range key or on an index return all values
             results = possible_results
 
         if index_name:
@@ -517,11 +525,17 @@ class DynamoDBBackend(BaseBackend):
                 for key in keys:
                     if key in table.hash_key_names:
                         return key, None
-
+            # import pdb; pdb.set_trace()
-            for potential_hash, potential_range in zip(table.hash_key_names, table.range_key_names):
+            # for potential_hash, potential_range in zip(table.hash_key_names, table.range_key_names):
-                if set([potential_hash, potential_range]) == set(keys):
+            #     if set([potential_hash, potential_range]) == set(keys):
-                    return potential_hash, potential_range
+            #         return potential_hash, potential_range
-            return None, None
+            potential_hash, potential_range = None, None
+            for key in set(keys):
+                if key in table.hash_key_names:
+                    potential_hash = key
+                elif key in table.range_key_names:
+                    potential_range = key
+            return potential_hash, potential_range
 
     def get_keys_value(self, table, keys):
         if table.hash_key_attr not in keys or (table.has_range_key and table.range_key_attr not in keys):
@@ -541,7 +555,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):
+              limit, exclusive_start_key, scan_index_forward, index_name=None, **filter_kwargs):
         table = self.tables.get(table_name)
         if not table:
             return None, None
@@ -550,7 +564,7 @@ class DynamoDBBackend(BaseBackend):
         range_values = [DynamoType(range_value) 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)
+                           exclusive_start_key, scan_index_forward, index_name, **filter_kwargs)
 
     def scan(self, table_name, filters, limit, exclusive_start_key):
         table = self.tables.get(table_name)

moto/dynamodb2/responses.py

@@ -240,6 +240,7 @@ 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')
+        filter_kwargs = {}
         if key_condition_expression:
             value_alias_map = self.body['ExpressionAttributeValues']
 
@@ -295,6 +296,9 @@ class DynamoHandler(BaseResponse):
             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())
+                for key, value in key_conditions.items():
+                    if key not in (hash_key_name, range_key_name):
+                        filter_kwargs[key] = value
                 if hash_key_name is None:
                     er = "'com.amazonaws.dynamodb.v20120810#ResourceNotFoundException"
                     return self.error(er)
@@ -303,25 +307,24 @@ class DynamoHandler(BaseResponse):
                     range_comparison = None
                     range_values = []
                 else:
-                    if range_key_name is None:
+                    if range_key_name is None and not filter_kwargs:
                         er = "com.amazon.coral.validate#ValidationException"
                         return self.error(er)
                     else:
-                        range_condition = key_conditions[range_key_name]
+                        range_condition = key_conditions.get(range_key_name)
                         if range_condition:
                             range_comparison = range_condition['ComparisonOperator']
                             range_values = range_condition['AttributeValueList']
                         else:
                             range_comparison = None
                             range_values = []
-
         index_name = self.body.get('IndexName')
         exclusive_start_key = self.body.get('ExclusiveStartKey')
         limit = self.body.get("Limit")
         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)
+            exclusive_start_key, scan_index_forward, index_name=index_name, **filter_kwargs)
         if items is None:
             er = 'com.amazonaws.dynamodb.v20111205#ResourceNotFoundException'
             return self.error(er)

tests/test_dynamodb2/test_dynamodb_table_with_range_key.py

@@ -653,11 +653,180 @@ def test_query_with_local_indexes():
 
     item['version'] = '2'
     item.save(overwrite=True)
-    # Revisit this query once support for QueryFilter is added
+    results = table.query(forum_name__eq='Cool Forum', index='threads_index', threads__eq=1)
-    results = table.query(forum_name__eq='Cool Forum', index='threads_index')
     list(results).should.have.length_of(1)
 
 
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_filter_eq():
+    table = create_table_with_local_indexes()
+    item_data = [
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Check this out!',
+            'version': '1',
+            'threads': 1,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Read this now!',
+            'version': '1',
+            'threads': 5,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Please read this... please',
+            'version': '1',
+            'threads': 0,
+        }
+    ]
+    for data in item_data:
+        item = Item(table, data)
+        item.save(overwrite=True)
+    results = table.query_2(
+        forum_name__eq='Cool Forum', index='threads_index', threads__eq=5
+    )
+    list(results).should.have.length_of(1)
+
+
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_filter_lt():
+    table = create_table_with_local_indexes()
+    item_data = [
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Check this out!',
+            'version': '1',
+            'threads': 1,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Read this now!',
+            'version': '1',
+            'threads': 5,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Please read this... please',
+            'version': '1',
+            'threads': 0,
+        }
+    ]
+    for data in item_data:
+        item = Item(table, data)
+        item.save(overwrite=True)
+
+    results = table.query(
+        forum_name__eq='Cool Forum', index='threads_index', threads__lt=5
+    )
+    results = list(results)
+    results.should.have.length_of(2)
+
+
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_filter_gt():
+    table = create_table_with_local_indexes()
+    item_data = [
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Check this out!',
+            'version': '1',
+            'threads': 1,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Read this now!',
+            'version': '1',
+            'threads': 5,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Please read this... please',
+            'version': '1',
+            'threads': 0,
+        }
+    ]
+    for data in item_data:
+        item = Item(table, data)
+        item.save(overwrite=True)
+
+    results = table.query(
+        forum_name__eq='Cool Forum', index='threads_index', threads__gt=1
+    )
+    list(results).should.have.length_of(1)
+
+
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_filter_lte():
+    table = create_table_with_local_indexes()
+    item_data = [
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Check this out!',
+            'version': '1',
+            'threads': 1,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Read this now!',
+            'version': '1',
+            'threads': 5,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Please read this... please',
+            'version': '1',
+            'threads': 0,
+        }
+    ]
+    for data in item_data:
+        item = Item(table, data)
+        item.save(overwrite=True)
+
+    results = table.query(
+        forum_name__eq='Cool Forum', index='threads_index', threads__lte=5
+    )
+    list(results).should.have.length_of(3)
+
+
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_filter_gte():
+    table = create_table_with_local_indexes()
+    item_data = [
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Check this out!',
+            'version': '1',
+            'threads': 1,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Read this now!',
+            'version': '1',
+            'threads': 5,
+        },
+        {
+            'forum_name': 'Cool Forum',
+            'subject': 'Please read this... please',
+            'version': '1',
+            'threads': 0,
+        }
+    ]
+    for data in item_data:
+        item = Item(table, data)
+        item.save(overwrite=True)
+
+    results = table.query(
+        forum_name__eq='Cool Forum', index='threads_index', threads__gte=1
+    )
+    list(results).should.have.length_of(2)
+
+
 @mock_dynamodb2
 def test_reverse_query():
     conn = boto.dynamodb2.layer1.DynamoDBConnection()