CRF MagicStack#295, use simple query ROLLBACK in implicit transaction only

Author: fantix

asyncpg/protocol/coreproto.pyx

@@ -810,34 +810,29 @@ cdef class CoreProtocol:
         self.transport.writelines(buffers)
 
     cdef _execute_many_rollback(self, object error):
-        cdef:
-            WriteBuffer packet
-            WriteBuffer buf
+        cdef WriteBuffer buf
 
         self._ensure_connected()
         if self.state != PROTOCOL_CANCELLED:
             self._set_state(PROTOCOL_BIND_EXECUTE_MANY_ROLLBACK)
-        packet = WriteBuffer.new()
-
-        # We have no idea if we are in an explicit transaction or not,
-        # therefore we raise an exception in the database to mark current
-        # transaction as failed anyway to prevent partial commit
-        buf = self._build_parse_message('',
-                                        'DO language plpgsql $$ BEGIN '
-                                        'RAISE transaction_rollback; '
-                                        'END$$;')
-        packet.write_buffer(buf)
-        buf = self._build_bind_message('', '', WriteBuffer.new())
-        packet.write_buffer(buf)
-        buf = self._build_execute_message('', 0)
-        packet.write_buffer(buf)
-        packet.write_bytes(SYNC_MESSAGE)
-
-        self.transport.write(memoryview(packet))
 
         self.result_type = RESULT_FAILED
         self.result = error
 
+        # We shall rollback in an implicit transaction to prevent partial
+        # commit, while do nothing in an explicit transaction and leaving the
+        # error to the user
+        if self.xact_status == PQTRANS_IDLE:
+            buf = WriteBuffer.new_message(b'Q')
+            # ROLLBACK here won't cause server to send RowDescription,
+            # CopyInResponse or CopyOutResponse which we are not expecting, but
+            # the server will send ReadyForQuery which finishes executemany()
+            buf.write_str('ROLLBACK;', self.encoding)
+            buf.end_message()
+            self._write(buf)
+        else:
+            self._write_sync_message()
+
     cdef _execute_many_end(self, object sync):
         if sync is True:
             self._write_sync_message()

tests/test_execute.py

@@ -231,3 +231,23 @@ async def test_execute_many_rollback(self):
             ''', AsyncIterable(sleep=0.3, sleep_start=32), timeout=0.5)
         result = await self.con.fetch('SELECT * FROM exmany')
         self.assertEqual(result, [])
+
+    async def test_execute_many_rollback_in_transaction(self):
+        # User error raised inside the transaction block
+        with self.assertRaises(ZeroDivisionError):
+            async with self.con.transaction():
+                await self.con.executemany('''
+                    INSERT INTO exmany VALUES($1, $2)
+                ''', (('a' * 32768, y + y / y) for y in range(10, -1, -1)))
+        result = await self.con.fetch('SELECT * FROM exmany')
+        self.assertEqual(result, [])
+
+        # User error captured inside the transaction block, executemany() is
+        # partially committed (4 * 2 buffers sent, [3, 2] not sent)
+        async with self.con.transaction():
+            with self.assertRaises(ZeroDivisionError):
+                await self.con.executemany('''
+                    INSERT INTO exmany VALUES($1, $2)
+                ''', (('a' * 32768, y + y / y) for y in range(10, -1, -1)))
+        result = await self.con.fetch('SELECT * FROM exmany')
+        self.assertEqual([x[1] for x in result], [11, 10, 9, 8, 7, 6, 5, 4])