Model.batch_get: add guard-rails (pynamodb#1184)

ikonst · web-flow · commit 12e127f5a821 · 2023-05-26T08:43:27.000-04:00
For models with a range key, fail if:
- item is a `str` ("accidental" iterable)
- item is an iterable with != 2 items
diff --git a/pynamodb/models.py b/pynamodb/models.py
@@ -357,15 +357,23 @@ def batch_get(
                         keys_to_get = []
             item = items.pop()
             if range_key_attribute:
-                hash_key, range_key = cls._serialize_keys(item[0], item[1])  # type: ignore
+                if isinstance(item, str):
+                    raise ValueError(f'Invalid key value {item!r}: '
+                                     'expected non-str iterable with exactly 2 elements (hash key, range key)')
+                try:
+                    hash_key, range_key = item
+                except (TypeError, ValueError):
+                    raise ValueError(f'Invalid key value {item!r}: '
+                                     'expected iterable with exactly 2 elements (hash key, range key)')
+                hash_key_ser, range_key_ser = cls._serialize_keys(hash_key, range_key)
                 keys_to_get.append({
-                    hash_key_attribute.attr_name: hash_key,
-                    range_key_attribute.attr_name: range_key
+                    hash_key_attribute.attr_name: hash_key_ser,
+                    range_key_attribute.attr_name: range_key_ser,
                 })
             else:
-                hash_key = cls._serialize_keys(item)[0]
+                hash_key_ser, _ = cls._serialize_keys(item)
                 keys_to_get.append({
-                    hash_key_attribute.attr_name: hash_key
+                    hash_key_attribute.attr_name: hash_key_ser
                 })
 
         while keys_to_get:
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -4,6 +4,7 @@
 import base64
 import json
 import copy
+import re
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
@@ -1845,7 +1846,6 @@ def test_batch_get(self):
             }
             self.assertEqual(params, req.call_args[0][1])
 
-
         with patch(PATCH_METHOD) as req:
             item_keys = [('hash-{}'.format(x), '{}'.format(x)) for x in range(10)]
             item_keys_copy = list(item_keys)
@@ -1906,6 +1906,56 @@ def fake_batch_get(*batch_args):
             for item in UserModel.batch_get(item_keys):
                 self.assertIsNotNone(item)
 
+    def test_batch_get__range_key(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = {
+                'UnprocessedKeys': {},
+                'Responses': {
+                    'UserModel': [],
+                }
+            }
+            items = [(f'hash-{x}', f'range-{x}') for x in range(10)]
+            _ = list(UserModel.batch_get(items))
+
+            actual_keys = req.call_args[0][1]['RequestItems']['UserModel']['Keys']
+            actual_keys.sort(key=json.dumps)
+            assert actual_keys == [
+                {'user_name': {'S': f'hash-{x}'}, 'user_id': {'S': f'range-{x}'}}
+                for x in range(10)
+            ]
+
+    def test_batch_get__range_key__invalid__string(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = {
+                'UnprocessedKeys': {},
+                'Responses': {
+                    'UserModel': [],
+                }
+            }
+            with pytest.raises(
+                ValueError,
+                match=re.escape(
+                    "Invalid key value 'ab': expected non-str iterable with exactly 2 elements (hash key, range key)"
+                )
+            ):
+                _ = list(UserModel.batch_get(['ab']))
+
+    def test_batch_get__range_key__invalid__3_elements(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = {
+                'UnprocessedKeys': {},
+                'Responses': {
+                    'UserModel': [],
+                }
+            }
+            with pytest.raises(
+                ValueError,
+                match=re.escape(
+                    "Invalid key value ('a', 'b', 'c'): expected iterable with exactly 2 elements (hash key, range key)"
+                )
+            ):
+                _ = list(UserModel.batch_get([('a', 'b', 'c')]))
+
     def test_batch_write(self):
         """
         Model.batch_write
