Handle exceptions in generators passed to Connection.executemany()

Fixes: #85
MagicStack · Mar 16, 2017 · 88ca460 · 88ca460
1 parent 71de129
commit 88ca460
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/asyncpg/protocol/coreproto.pyx b/asyncpg/protocol/coreproto.pyx
@@ -235,6 +235,10 @@ cdef class CoreProtocol:
                     buf = <WriteBuffer>next(self._execute_iter)
                 except StopIteration:
                     self._push_result()
+                except Exception as e:
+                    self.result_type = RESULT_FAILED
+                    self.result = e
+                    self._push_result()
                 else:
                     # Next iteration over the executemany() arg sequence
                     self._send_bind_message(
@@ -643,6 +647,10 @@ cdef class CoreProtocol:
             buf = <WriteBuffer>next(bind_data)
         except StopIteration:
             self._push_result()
+        except Exception as e:
+            self.result_type = RESULT_FAILED
+            self.result = e
+            self._push_result()
         else:
             self._send_bind_message(portal_name, stmt_name, buf, 0)
 

diff --git a/tests/test_execute.py b/tests/test_execute.py
@@ -132,3 +132,23 @@ async def test_execute_many_1(self):
             ])
         finally:
             await self.con.execute('DROP TABLE exmany')
+
+    async def test_execute_many_2(self):
+        await self.con.execute('CREATE TEMP TABLE exmany (b int)')
+
+        try:
+            bad_data = ([1 / 0] for v in range(10))
+
+            with self.assertRaises(ZeroDivisionError):
+                async with self.con.transaction():
+                    await self.con.executemany('''
+                        INSERT INTO exmany VALUES($1)
+                    ''', bad_data)
+
+            good_data = ([v] for v in range(10))
+            async with self.con.transaction():
+                await self.con.executemany('''
+                    INSERT INTO exmany VALUES($1)
+                ''', good_data)
+        finally:
+            await self.con.execute('DROP TABLE exmany')