introduce concept of page_size, separate from num items returned limit (#139)

Author: anandswaminathan

pynamodb/models.py

@@ -510,6 +510,7 @@ def query(cls,
               limit=None,
               last_evaluated_key=None,
               attributes_to_get=None,
+              page_size=None,
               **filters):
         """
         Provides a high level query API
@@ -522,6 +523,7 @@ def query(cls,
             Controls descending or ascending results
         :param last_evaluated_key: If set, provides the starting point for query.
         :param attributes_to_get: If set, only returns these elements
+        :param page_size: Page size of the query to DynamoDB
         :param filters: A dictionary of filters to be used in the query
         """
         cls._get_indexes()
@@ -538,6 +540,10 @@ def query(cls,
                     key_attribute_classes[name] = attr
                 else:
                     non_key_attribute_classes[name] = attr
+
+        if page_size is None:
+            page_size = limit
+
         key_conditions, query_filters = cls._build_filters(
             QUERY_OPERATOR_MAP,
             non_key_operator_map=QUERY_FILTER_OPERATOR_MAP,
@@ -551,7 +557,7 @@ def query(cls,
             exclusive_start_key=last_evaluated_key,
             consistent_read=consistent_read,
             scan_index_forward=scan_index_forward,
-            limit=limit,
+            limit=page_size,
             key_conditions=key_conditions,
             attributes_to_get=attributes_to_get,
             query_filters=query_filters,
@@ -571,14 +577,7 @@ def query(cls,
             yield cls.from_raw_data(item)
 
         while last_evaluated_key:
-            log.debug("Fetching query page with exclusive start key: %s", last_evaluated_key)
-            # If the user provided a limit, we need to subtract the number of results returned for each page
-            if limit is not None:
-                if limit == 0:
-                    return
-                limit -= data.get(CAMEL_COUNT, 0)
             query_kwargs['exclusive_start_key'] = last_evaluated_key
-            query_kwargs['limit'] = limit
             log.debug("Fetching query page with exclusive start key: %s", last_evaluated_key)
             data = cls._get_connection().query(hash_key, **query_kwargs)
             cls._throttle.add_record(data.get(CONSUMED_CAPACITY))
@@ -597,6 +596,7 @@ def scan(cls,
              limit=None,
              conditional_operator=None,
              last_evaluated_key=None,
+             page_size=None,
              **filters):
         """
         Iterates through all items in the table
@@ -605,6 +605,7 @@ def scan(cls,
         :param total_segments: If set, then specifies total segments
         :param limit: Used to limit the number of results returned
         :param last_evaluated_key: If set, provides the starting point for scan.
+        :param page_size: Page size of the scan to DynamoDB
         :param filters: A list of item filters
         """
         key_filter, scan_filter = cls._build_filters(
@@ -614,10 +615,13 @@ def scan(cls,
             filters=filters
         )
         key_filter.update(scan_filter)
+        if page_size is None:
+            page_size = limit
+
         data = cls._get_connection().scan(
             exclusive_start_key=last_evaluated_key,
             segment=segment,
-            limit=limit,
+            limit=page_size,
             scan_filter=key_filter,
             total_segments=total_segments,
             conditional_operator=conditional_operator
@@ -635,7 +639,7 @@ def scan(cls,
             log.debug("Fetching scan page with exclusive start key: %s", last_evaluated_key)
             data = cls._get_connection().scan(
                 exclusive_start_key=last_evaluated_key,
-                limit=limit,
+                limit=page_size,
                 scan_filter=key_filter,
                 segment=segment,
                 total_segments=total_segments

pynamodb/tests/test_model.py

@@ -1174,8 +1174,32 @@ def test_query_limit_less_than_available_items_multiple_page(self):
             self.assertEqual(len(results), 25)
             self.assertEqual(len(req.mock_calls), 3)
             self.assertEquals(req.mock_calls[0][1][1]['Limit'], 25)
-            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 15)
-            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 5)
+            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 25)
+            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 25)
+
+    def test_query_limit_less_than_available_and_page_size(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = MODEL_TABLE_DATA
+            UserModel('foo', 'bar')
+
+        with patch(PATCH_METHOD) as req:
+            items = []
+            for idx in range(30):
+                item = copy.copy(GET_MODEL_ITEM_DATA.get(ITEM))
+                item['user_id'] = {STRING_SHORT: 'id-{0}'.format(idx)}
+                items.append(item)
+
+            req.side_effect = [
+                {'Items': items[:10], 'LastEvaluatedKey': 'x'},
+                {'Items': items[10:20], 'LastEvaluatedKey': 'y'},
+                {'Items': items[20:30]},
+            ]
+            results = list(UserModel.query('foo', limit=25, page_size=10))
+            self.assertEqual(len(results), 25)
+            self.assertEqual(len(req.mock_calls), 3)
+            self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 10)
 
     def test_query_limit_greater_than_available_items_multiple_page(self):
         with patch(PATCH_METHOD) as req:
@@ -1198,8 +1222,32 @@ def test_query_limit_greater_than_available_items_multiple_page(self):
             self.assertEqual(len(results), 30)
             self.assertEqual(len(req.mock_calls), 3)
             self.assertEquals(req.mock_calls[0][1][1]['Limit'], 50)
-            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 40)
-            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 30)
+            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 50)
+            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 50)
+
+    def test_query_limit_greater_than_available_items_and_page_size(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = MODEL_TABLE_DATA
+            UserModel('foo', 'bar')
+
+        with patch(PATCH_METHOD) as req:
+            items = []
+            for idx in range(30):
+                item = copy.copy(GET_MODEL_ITEM_DATA.get(ITEM))
+                item['user_id'] = {STRING_SHORT: 'id-{0}'.format(idx)}
+                items.append(item)
+
+            req.side_effect = [
+                {'Items': items[:10], 'LastEvaluatedKey': 'x'},
+                {'Items': items[10:20], 'LastEvaluatedKey': 'y'},
+                {'Items': items[20:30]},
+            ]
+            results = list(UserModel.query('foo', limit=50, page_size=10))
+            self.assertEqual(len(results), 30)
+            self.assertEqual(len(req.mock_calls), 3)
+            self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 10)
 
     def test_query(self):
         """
@@ -1390,6 +1438,26 @@ def fake_query(*args):
             self.assertEqual(params, req.call_args[0][1])
             self.assertTrue(len(queried) == len(items))
 
+    def test_scan_limit_with_page_size(self):
+        with patch(PATCH_METHOD) as req:
+            items = []
+            for idx in range(30):
+                item = copy.copy(GET_MODEL_ITEM_DATA.get(ITEM))
+                item['user_id'] = {STRING_SHORT: 'id-{0}'.format(idx)}
+                items.append(item)
+
+            req.side_effect = [
+                {'Items': items[:10], 'LastEvaluatedKey': 'x'},
+                {'Items': items[10:20], 'LastEvaluatedKey': 'y'},
+                {'Items': items[20:30]},
+            ]
+            results = list(UserModel.scan(limit=25, page_size=10))
+            self.assertEqual(len(results), 25)
+            self.assertEqual(len(req.mock_calls), 3)
+            self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[1][1][1]['Limit'], 10)
+            self.assertEquals(req.mock_calls[2][1][1]['Limit'], 10)
+
     def test_scan_limit(self):
         """
         Model.scan(limit)
@@ -1411,6 +1479,8 @@ def fake_scan(*args):
             for item in UserModel.scan(limit=4):
                 count += 1
                 self.assertIsNotNone(item)
+            self.assertEqual(len(req.mock_calls), 1)
+            self.assertEquals(req.mock_calls[0][1][1]['Limit'], 4)
             self.assertEqual(count, 4)
 
         with patch(PATCH_METHOD) as req: