Query and Scan Filters using Condition Expressions (#339)

Author: jpinner-lyft

pynamodb/connection/base.py

@@ -109,6 +109,24 @@ def get_index_hash_keyname(self, index_name):
                     if schema_key.get(KEY_TYPE) == HASH:
                         return schema_key.get(ATTR_NAME)
 
+    def get_index_range_keyname(self, index_name):
+        """
+        Returns the name of the hash key for a given index
+        """
+        global_indexes = self.data.get(GLOBAL_SECONDARY_INDEXES)
+        local_indexes = self.data.get(LOCAL_SECONDARY_INDEXES)
+        indexes = []
+        if local_indexes:
+            indexes += local_indexes
+        if global_indexes:
+            indexes += global_indexes
+        for index in indexes:
+            if index.get(INDEX_NAME) == index_name:
+                for schema_key in index.get(KEY_SCHEMA):
+                    if schema_key.get(KEY_TYPE) == RANGE:
+                        return schema_key.get(ATTR_NAME)
+        return None
+
     def get_item_attribute_map(self, attributes, item_key=ITEM, pythonic_key=True):
         """
         Builds up a dynamodb compatible AttributeValue map
@@ -1031,6 +1049,7 @@ def get_item(self,
 
     def rate_limited_scan(self,
                           table_name,
+                          filter_condition=None,
                           attributes_to_get=None,
                           page_size=None,
                           limit=None,
@@ -1051,6 +1070,7 @@ def rate_limited_scan(self,
         limit the rate of the scan. 'ProvisionedThroughputExceededException' is also handled and retried.
 
         :param table_name: Name of the table to perform scan on.
+        :param filter_condition: Condition used to restrict the scan results
         :param attributes_to_get: A list of attributes to return.
         :param page_size: Page size of the scan to DynamoDB
         :param limit: Used to limit the number of results returned
@@ -1094,6 +1114,7 @@ def rate_limited_scan(self,
                 try:
                     data = self.scan(
                         table_name,
+                        filter_condition=filter_condition,
                         attributes_to_get=attributes_to_get,
                         exclusive_start_key=last_evaluated_key,
                         limit=page_size,
@@ -1177,6 +1198,7 @@ def rate_limited_scan(self,
 
     def scan(self,
              table_name,
+             filter_condition=None,
              attributes_to_get=None,
              limit=None,
              conditional_operator=None,
@@ -1189,10 +1211,15 @@ def scan(self,
         """
         Performs the scan operation
         """
+        self._check_condition('filter_condition', filter_condition, scan_filter, conditional_operator)
+
         operation_kwargs = {TABLE_NAME: table_name}
         name_placeholders = {}
         expression_attribute_values = {}
 
+        if filter_condition is not None:
+            filter_expression = filter_condition.serialize(name_placeholders, expression_attribute_values)
+            operation_kwargs[FILTER_EXPRESSION] = filter_expression
         if attributes_to_get is not None:
             projection_expression = create_projection_expression(attributes_to_get, name_placeholders)
             operation_kwargs[PROJECTION_EXPRESSION] = projection_expression
@@ -1226,6 +1253,8 @@ def scan(self,
     def query(self,
               table_name,
               hash_key,
+              range_key_condition=None,
+              filter_condition=None,
               attributes_to_get=None,
               consistent_read=False,
               exclusive_start_key=None,
@@ -1240,6 +1269,9 @@ def query(self,
         """
         Performs the Query operation and returns the result
         """
+        self._check_condition('range_key_condition', range_key_condition, key_conditions, conditional_operator)
+        self._check_condition('filter_condition', filter_condition, query_filters, conditional_operator)
+
         operation_kwargs = {TABLE_NAME: table_name}
         name_placeholders = {}
         expression_attribute_values = {}
@@ -1251,10 +1283,21 @@ def query(self,
             hash_keyname = tbl.get_index_hash_keyname(index_name)
             if not hash_keyname:
                 raise ValueError("No hash key attribute for index: {0}".format(index_name))
+            range_keyname = tbl.get_index_range_keyname(index_name)
         else:
             hash_keyname = tbl.hash_keyname
+            range_keyname = tbl.range_keyname
+
+        key_condition = self._get_condition(table_name, hash_keyname, '__eq__', hash_key)
+        if range_key_condition is not None:
+            if range_key_condition.is_valid_range_key_condition(range_keyname):
+                key_condition = key_condition & range_key_condition
+            elif filter_condition is None:
+                # Try to gracefully handle the case where a user passed in a filter as a range key condition
+                (filter_condition, range_key_condition) = (range_key_condition, None)
+            else:
+                raise ValueError("{0} is not a valid range key condition".format(range_key_condition))
 
-        key_condition_expression = self._get_condition(table_name, hash_keyname, '__eq__', hash_key)
         if key_conditions is None or len(key_conditions) == 0:
             pass  # No comparisons on sort key
         elif len(key_conditions) > 1:
@@ -1267,11 +1310,21 @@ def query(self,
             operator = KEY_CONDITION_OPERATOR_MAP[operator]
             values = condition.get(ATTR_VALUE_LIST)
             sort_key_expression = self._get_condition(table_name, key, operator, *values)
-            key_condition_expression = key_condition_expression & sort_key_expression
+            key_condition = key_condition & sort_key_expression
 
-        operation_kwargs[KEY_CONDITION_EXPRESSION] = key_condition_expression.serialize(
+        operation_kwargs[KEY_CONDITION_EXPRESSION] = key_condition.serialize(
             name_placeholders, expression_attribute_values)
-
+        if filter_condition is not None:
+            filter_expression = filter_condition.serialize(name_placeholders, expression_attribute_values)
+            # FilterExpression does not allow key attributes. Check for hash and range key name placeholders
+            hash_key_placeholder = name_placeholders.get(hash_keyname)
+            range_key_placeholder = range_keyname and name_placeholders.get(range_keyname)
+            if (
+                hash_key_placeholder in filter_expression or
+                (range_key_placeholder and range_key_placeholder in filter_expression)
+            ):
+                raise ValueError("'filter_condition' cannot contain key attributes")
+            operation_kwargs[FILTER_EXPRESSION] = filter_expression
         if attributes_to_get:
             projection_expression = create_projection_expression(attributes_to_get, name_placeholders)
             operation_kwargs[PROJECTION_EXPRESSION] = projection_expression
@@ -1383,7 +1436,7 @@ def _check_condition(self, name, condition, expected_or_filter, conditional_oper
         if condition is not None:
             if not isinstance(condition, Condition):
                 raise ValueError("'{0}' must be an instance of Condition".format(name))
-            if expected_or_filter is not None or conditional_operator is not None:
+            if expected_or_filter or conditional_operator is not None:
                 raise ValueError("Legacy conditional parameters cannot be used with condition expressions")
 
     @staticmethod

pynamodb/connection/table.py

@@ -138,6 +138,7 @@ def get_item(self, hash_key, range_key=None, consistent_read=False, attributes_t
             attributes_to_get=attributes_to_get)
 
     def rate_limited_scan(self,
+             filter_condition=None,
              attributes_to_get=None,
              page_size=None,
              limit=None,
@@ -157,6 +158,7 @@ def rate_limited_scan(self,
         """
         return self.connection.rate_limited_scan(
             self.table_name,
+            filter_condition=filter_condition,
             attributes_to_get=attributes_to_get,
             page_size=page_size,
             limit=limit,
@@ -173,6 +175,7 @@ def rate_limited_scan(self,
             consistent_read=consistent_read)
 
     def scan(self,
+             filter_condition=None,
              attributes_to_get=None,
              limit=None,
              conditional_operator=None,
@@ -187,6 +190,7 @@ def scan(self,
         """
         return self.connection.scan(
             self.table_name,
+            filter_condition=filter_condition,
             attributes_to_get=attributes_to_get,
             limit=limit,
             conditional_operator=conditional_operator,
@@ -199,6 +203,8 @@ def scan(self,
 
     def query(self,
               hash_key,
+              range_key_condition=None,
+              filter_condition=None,
               attributes_to_get=None,
               consistent_read=False,
               exclusive_start_key=None,
@@ -217,6 +223,8 @@ def query(self,
         return self.connection.query(
             self.table_name,
             hash_key,
+            range_key_condition=range_key_condition,
+            filter_condition=filter_condition,
             attributes_to_get=attributes_to_get,
             consistent_read=consistent_read,
             exclusive_start_key=exclusive_start_key,

pynamodb/expressions/condition.py

@@ -137,6 +137,10 @@ def __init__(self, path, operator, *values):
         self.operator = operator
         self.values = values
 
+    # http://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_Query.html#DDB-Query-request-KeyConditionExpression
+    def is_valid_range_key_condition(self, path):
+        return str(self.path) == path and self.operator in ['=', '<', '<=', '>', '>=', BETWEEN, 'begins_with']
+
     def serialize(self, placeholder_names, expression_attribute_values):
         path = self._get_path(self.path, placeholder_names)
         values = self._get_values(placeholder_names, expression_attribute_values)
@@ -175,7 +179,7 @@ def __invert__(self):
         return Not(self)
 
     def __repr__(self):
-        values = [repr(value) if isinstance(value, Condition) else value.items()[0][1] for value in self.values]
+        values = [repr(value) if isinstance(value, Condition) else list(value.items())[0][1] for value in self.values]
         return self.format_string.format(*values, path=self.path, operator = self.operator)
 
     def __nonzero__(self):

pynamodb/indexes.py

@@ -45,13 +45,17 @@ def __init__(self):
     @classmethod
     def count(cls,
               hash_key,
+              range_key_condition=None,
+              filter_condition=None,
               consistent_read=False,
               **filters):
         """
         Count on an index
         """
         return cls.Meta.model.count(
             hash_key,
+            range_key_condition=range_key_condition,
+            filter_condition=filter_condition,
             index_name=cls.Meta.index_name,
             consistent_read=consistent_read,
             **filters
@@ -60,6 +64,8 @@ def count(cls,
     @classmethod
     def query(self,
               hash_key,
+              range_key_condition=None,
+              filter_condition=None,
               scan_index_forward=None,
               consistent_read=False,
               limit=None,
@@ -71,6 +77,8 @@ def query(self,
         """
         return self.Meta.model.query(
             hash_key,
+            range_key_condition=range_key_condition,
+            filter_condition=filter_condition,
             index_name=self.Meta.index_name,
             scan_index_forward=scan_index_forward,
             consistent_read=consistent_read,

pynamodb/models.py

@@ -517,6 +517,8 @@ def from_raw_data(cls, data):
     @classmethod
     def count(cls,
               hash_key=None,
+              range_key_condition=None,
+              filter_condition=None,
               consistent_read=False,
               index_name=None,
               limit=None,
@@ -525,6 +527,8 @@ def count(cls,
         Provides a filtered count
 
         :param hash_key: The hash key to query. Can be None.
+        :param range_key_condition: Condition for range key
+        :param filter_condition: Condition used to restrict the query results
         :param consistent_read: If True, a consistent read is performed
         :param index_name: If set, then this index is used
         :param filters: A dictionary of filters to be used in the query. Requires a hash_key to be passed.
@@ -562,6 +566,8 @@ def count(cls,
             started = True
             data = cls._get_connection().query(
                 hash_key,
+                range_key_condition=range_key_condition,
+                filter_condition=filter_condition,
                 index_name=index_name,
                 consistent_read=consistent_read,
                 key_conditions=key_conditions,
@@ -578,6 +584,8 @@ def count(cls,
     @classmethod
     def query(cls,
               hash_key,
+              range_key_condition=None,
+              filter_condition=None,
               consistent_read=False,
               index_name=None,
               scan_index_forward=None,
@@ -591,6 +599,8 @@ def query(cls,
         Provides a high level query API
 
         :param hash_key: The hash key to query
+        :param range_key_condition: Condition for range key
+        :param filter_condition: Condition used to restrict the query results
         :param consistent_read: If True, a consistent read is performed
         :param index_name: If set, then this index is used
         :param limit: Used to limit the number of results returned
@@ -630,6 +640,8 @@ def query(cls,
         log.debug("Fetching first query page")
 
         query_kwargs = dict(
+            range_key_condition=range_key_condition,
+            filter_condition=filter_condition,
             index_name=index_name,
             exclusive_start_key=last_evaluated_key,
             consistent_read=consistent_read,
@@ -668,6 +680,7 @@ def query(cls,
 
     @classmethod
     def rate_limited_scan(cls,
+            filter_condition=None,
             attributes_to_get=None,
             segment=None,
             total_segments=None,
@@ -686,6 +699,7 @@ def rate_limited_scan(cls,
         Scans the items in the table at a definite rate.
         Invokes the low level rate_limited_scan API.
 
+        :param filter_condition: Condition used to restrict the scan results
         :param attributes_to_get: A list of attributes to return.
         :param segment: If set, then scans the segment
         :param total_segments: If set, then specifies total segments
@@ -717,6 +731,7 @@ def rate_limited_scan(cls,
         key_filter.update(scan_filter)
 
         scan_result = cls._get_connection().rate_limited_scan(
+            filter_condition=filter_condition,
             attributes_to_get=attributes_to_get,
             page_size=page_size,
             limit=limit,
@@ -738,6 +753,7 @@ def rate_limited_scan(cls,
 
     @classmethod
     def scan(cls,
+             filter_condition=None,
              segment=None,
              total_segments=None,
              limit=None,
@@ -749,6 +765,7 @@ def scan(cls,
         """
         Iterates through all items in the table
 
+        :param filter_condition: Condition used to restrict the scan results
         :param segment: If set, then scans the segment
         :param total_segments: If set, then specifies total segments
         :param limit: Used to limit the number of results returned
@@ -770,6 +787,7 @@ def scan(cls,
             page_size = limit
 
         data = cls._get_connection().scan(
+            filter_condition=filter_condition,
             exclusive_start_key=last_evaluated_key,
             segment=segment,
             limit=page_size,
@@ -790,6 +808,7 @@ def scan(cls,
         while last_evaluated_key:
             log.debug("Fetching scan page with exclusive start key: %s", last_evaluated_key)
             data = cls._get_connection().scan(
+                filter_condition=filter_condition,
                 exclusive_start_key=last_evaluated_key,
                 limit=page_size,
                 scan_filter=key_filter,