pass batch parameters down to boto/dynamo (#308)

Author: jmphilli

pynamodb/models.py

@@ -260,8 +260,11 @@ def batch_get(cls, items, consistent_read=None, attributes_to_get=None):
         while items:
             if len(keys_to_get) == BATCH_GET_PAGE_LIMIT:
                 while keys_to_get:
-                    page, unprocessed_keys = cls._batch_get_page(keys_to_get, consistent_read=None,
-                                                                 attributes_to_get=None)
+                    page, unprocessed_keys = cls._batch_get_page(
+                        keys_to_get,
+                        consistent_read=consistent_read,
+                        attributes_to_get=attributes_to_get
+                    )
                     for batch_item in page:
                         yield cls.from_raw_data(batch_item)
                     if unprocessed_keys:
@@ -282,7 +285,11 @@ def batch_get(cls, items, consistent_read=None, attributes_to_get=None):
                 })
 
         while keys_to_get:
-            page, unprocessed_keys = cls._batch_get_page(keys_to_get, consistent_read=None, attributes_to_get=None)
+            page, unprocessed_keys = cls._batch_get_page(
+                keys_to_get,
+                consistent_read=consistent_read,
+                attributes_to_get=attributes_to_get
+            )
             for batch_item in page:
                 yield cls.from_raw_data(batch_item)
             if unprocessed_keys:
@@ -1209,7 +1216,7 @@ def _batch_get_page(cls, keys_to_get, consistent_read, attributes_to_get):
         """
         log.debug("Fetching a BatchGetItem page")
         data = cls._get_connection().batch_get_item(
-            keys_to_get, consistent_read, attributes_to_get
+            keys_to_get, consistent_read=consistent_read, attributes_to_get=attributes_to_get
         )
         cls._throttle.add_record(data.get(CONSUMED_CAPACITY))
         item_data = data.get(RESPONSES).get(cls.Meta.table_name)

pynamodb/tests/test_model.py

@@ -2344,6 +2344,33 @@ def test_batch_get(self):
             }
             self.assertEqual(params, req.call_args[0][1])
 
+        with patch(PATCH_METHOD) as req:
+            req.return_value = SIMPLE_BATCH_GET_ITEMS
+            item_keys = ['hash-{0}'.format(x) for x in range(10)]
+            for item in SimpleUserModel.batch_get(item_keys, attributes_to_get=['numbers']):
+                self.assertIsNotNone(item)
+            params = {
+                'ReturnConsumedCapacity': 'TOTAL',
+                'RequestItems': {
+                    'SimpleModel': {
+                        'Keys': [
+                            {'user_name': {'S': 'hash-9'}},
+                            {'user_name': {'S': 'hash-8'}},
+                            {'user_name': {'S': 'hash-7'}},
+                            {'user_name': {'S': 'hash-6'}},
+                            {'user_name': {'S': 'hash-5'}},
+                            {'user_name': {'S': 'hash-4'}},
+                            {'user_name': {'S': 'hash-3'}},
+                            {'user_name': {'S': 'hash-2'}},
+                            {'user_name': {'S': 'hash-1'}},
+                            {'user_name': {'S': 'hash-0'}}
+                        ],
+                        'AttributesToGet': ['numbers']
+                    }
+                }
+            }
+            self.assertEqual(params, req.call_args[0][1])
+
         with patch(PATCH_METHOD) as req:
             req.return_value = MODEL_TABLE_DATA
             UserModel('foo', 'bar')