Fix StopIteration in deduplication rows code when conflict_action == ConflictAction.NOTHING and rows parameter is iterator or generator

Author: seroy

psqlextra/query.py

@@ -174,11 +174,19 @@ def bulk_insert(
             A list of either the dicts of the rows inserted, including the pk or
             the models of the rows inserted with defaults for any fields not specified
         """
+        if rows is None:
+            return []
+
+        def peek(iterable):
+            try:
+                first = next(iterable)
+            except StopIteration:
+                return None
+            return list(chain([first], iterable))
 
-        def is_empty(r):
-            return all([False for _ in r])
+        rows = peek(iter(rows))
 
-        if not rows or is_empty(rows):
+        if not rows:
             return []
 
         if not self.conflict_target and not self.conflict_action:

tests/test_on_conflict_nothing.py

@@ -179,8 +179,15 @@ def test_on_conflict_nothing_duplicate_rows():
 
     rows = [dict(amount=1), dict(amount=1)]
 
-    (
-        model.objects.on_conflict(
-            ["amount"], ConflictAction.NOTHING
-        ).bulk_insert(rows)
-    )
+    inserted_rows = model.objects.on_conflict(
+        ["amount"], ConflictAction.NOTHING
+    ).bulk_insert(rows)
+
+    assert len(inserted_rows) == 1
+
+    rows = iter([dict(amount=2), dict(amount=2)])
+    inserted_rows = model.objects.on_conflict(
+        ["amount"], ConflictAction.NOTHING
+    ).bulk_insert(rows)
+
+    assert len(inserted_rows) == 1