Rewrite again executemany() in a much readable way.

Now `Bind` and `Execute` pairs are batched into 4 x 32KB buffers to take
advantage of `writelines()`. A single `Sync` is sent at last, so that
all args live in the same transaction.

Closes: MagicStack#289

Author: fantix

asyncpg/connection.py

@@ -301,6 +301,13 @@ async def executemany(self, command: str, args, *, timeout: float=None):
 
         .. versionchanged:: 0.11.0
            `timeout` became a keyword-only parameter.
+
+        .. versionchanged:: 0.19.0
+           The execution was changed to be in an implicit transaction if there
+           was no explicit transaction, so that it will no longer end up with
+           partial success. If you still need the previous behavior to
+           progressively execute many args, please use a loop with prepared
+           statement instead.
         """
         self._check_open()
         return await self._executemany(command, args, timeout)

asyncpg/prepared_stmt.py

@@ -196,11 +196,24 @@ async def fetchrow(self, *args, timeout=None):
             return None
         return data[0]
 
-    async def __bind_execute(self, args, limit, timeout):
+    @connresource.guarded
+    async def executemany(self, args, *, timeout: float=None):
+        """Execute the statement for each sequence of arguments in *args*.
+
+        :param args: An iterable containing sequences of arguments.
+        :param float timeout: Optional timeout value in seconds.
+        :return None: This method discards the results of the operations.
+
+        .. versionadded:: 0.19.0
+        """
+        return await self.__do_execute(
+            lambda protocol: protocol.bind_execute_many(
+                self._state, args, '', timeout))
+
+    async def __do_execute(self, executor):
         protocol = self._connection._protocol
         try:
-            data, status, _ = await protocol.bind_execute(
-                self._state, args, '', limit, True, timeout)
+            return await executor(protocol)
         except exceptions.OutdatedSchemaCacheError:
             await self._connection.reload_schema_state()
             # We can not find all manually created prepared statements, so just
@@ -209,6 +222,11 @@ async def __bind_execute(self, args, limit, timeout):
             # invalidate themselves (unfortunately, clearing caches again).
             self._state.mark_closed()
             raise
+
+    async def __bind_execute(self, args, limit, timeout):
+        data, status, _ = await self.__do_execute(
+            lambda protocol: protocol.bind_execute(
+                self._state, args, '', limit, True, timeout))
         self._last_status = status
         return data
 

asyncpg/protocol/consts.pxi

@@ -8,3 +8,5 @@
 DEF _MAXINT32 = 2**31 - 1
 DEF _COPY_BUFFER_SIZE = 524288
 DEF _COPY_SIGNATURE = b"PGCOPY\n\377\r\n\0"
+DEF _EXECUTE_MANY_BUF_NUM = 4
+DEF _EXECUTE_MANY_BUF_SIZE = 32768

asyncpg/protocol/coreproto.pxd

@@ -76,11 +76,6 @@ cdef class CoreProtocol:
         bint _skip_discard
         bint _discard_data
 
-        # executemany support data
-        object _execute_iter
-        str _execute_portal_name
-        str _execute_stmt_name
-
         ConnectionStatus con_status
         ProtocolState state
         TransactionStatus xact_status
@@ -105,6 +100,7 @@ cdef class CoreProtocol:
         # True - completed, False - suspended
         bint result_execute_completed
 
+    cpdef is_in_transaction(self)
     cdef _process__auth(self, char mtype)
     cdef _process__prepare(self, char mtype)
     cdef _process__bind_execute(self, char mtype)
@@ -135,6 +131,7 @@ cdef class CoreProtocol:
     cdef _auth_password_message_md5(self, bytes salt)
 
     cdef _write(self, buf)
+    cdef _writelines(self, list buffers)
 
     cdef _read_server_messages(self)
 
@@ -147,6 +144,8 @@ cdef class CoreProtocol:
     cdef WriteBuffer _build_bind_message(self, str portal_name,
                                          str stmt_name,
                                          WriteBuffer bind_data)
+    cdef WriteBuffer _build_execute_message(self, str portal_name,
+                                            int32_t limit)
 
 
     cdef _connect(self)
@@ -155,8 +154,11 @@ cdef class CoreProtocol:
                             WriteBuffer bind_data, int32_t limit)
     cdef _bind_execute(self, str portal_name, str stmt_name,
                        WriteBuffer bind_data, int32_t limit)
-    cdef _bind_execute_many(self, str portal_name, str stmt_name,
-                            object bind_data)
+    cdef _execute_many_init(self)
+    cdef _execute_many_writelines(self, str portal_name, str stmt_name,
+                                  object bind_data)
+    cdef _execute_many_done(self, bint data_sent)
+    cdef _execute_many_fail(self, object error)
     cdef _bind(self, str portal_name, str stmt_name,
                WriteBuffer bind_data)
     cdef _execute(self, str portal_name, int32_t limit)

asyncpg/protocol/coreproto.pyx

@@ -22,13 +22,13 @@ cdef class CoreProtocol:
         self.xact_status = PQTRANS_IDLE
         self.encoding = 'utf-8'
 
-        # executemany support data
-        self._execute_iter = None
-        self._execute_portal_name = None
-        self._execute_stmt_name = None
-
         self._reset_result()
 
+    cpdef is_in_transaction(self):
+        # PQTRANS_INTRANS = idle, within transaction block
+        # PQTRANS_INERROR = idle, within failed transaction
+        return self.xact_status in (PQTRANS_INTRANS, PQTRANS_INERROR)
+
     cdef _read_server_messages(self):
         cdef:
             char mtype
@@ -253,22 +253,7 @@ cdef class CoreProtocol:
         elif mtype == b'Z':
             # ReadyForQuery
             self._parse_msg_ready_for_query()
-            if self.result_type == RESULT_FAILED:
-                self._push_result()
-            else:
-                try:
-                    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(
-                        self._execute_portal_name, self._execute_stmt_name,
-                        buf, 0)
+            self._push_result()
 
         elif mtype == b'I':
             # EmptyQueryResponse
@@ -702,6 +687,17 @@ cdef class CoreProtocol:
         buf.end_message()
         return buf
 
+    cdef WriteBuffer _build_execute_message(self, str portal_name,
+                                            int32_t limit):
+        cdef WriteBuffer buf
+
+        buf = WriteBuffer.new_message(b'E')
+        buf.write_str(portal_name, self.encoding)  # name of the portal
+        buf.write_int32(limit)  # number of rows to return; 0 - all
+
+        buf.end_message()
+        return buf
+
     # API for subclasses
 
     cdef _connect(self):
@@ -779,10 +775,7 @@ cdef class CoreProtocol:
         buf = self._build_bind_message(portal_name, stmt_name, bind_data)
         packet = buf
 
-        buf = WriteBuffer.new_message(b'E')
-        buf.write_str(portal_name, self.encoding)  # name of the portal
-        buf.write_int32(limit)  # number of rows to return; 0 - all
-        buf.end_message()
+        buf = self._build_execute_message(portal_name, limit)
         packet.write_buffer(buf)
 
         packet.write_bytes(SYNC_MESSAGE)
@@ -801,30 +794,71 @@ cdef class CoreProtocol:
 
         self._send_bind_message(portal_name, stmt_name, bind_data, limit)
 
-    cdef _bind_execute_many(self, str portal_name, str stmt_name,
-                            object bind_data):
-
-        cdef WriteBuffer buf
-
+    cdef _execute_many_init(self):
         self._ensure_connected()
         self._set_state(PROTOCOL_BIND_EXECUTE_MANY)
 
         self.result = None
         self._discard_data = True
-        self._execute_iter = bind_data
-        self._execute_portal_name = portal_name
-        self._execute_stmt_name = stmt_name
 
-        try:
-            buf = <WriteBuffer>next(bind_data)
-        except StopIteration:
-            self._push_result()
-        except Exception as e:
-            self.result_type = RESULT_FAILED
-            self.result = e
+    cdef _execute_many_writelines(self, str portal_name, str stmt_name,
+                                  object bind_data):
+        cdef:
+            WriteBuffer packet
+            WriteBuffer buf
+            list buffers = []
+
+        if self.result_type == RESULT_FAILED:
+            raise StopIteration(False)
+
+        while len(buffers) < _EXECUTE_MANY_BUF_NUM:
+            packet = WriteBuffer.new()
+
+            while packet.len() < _EXECUTE_MANY_BUF_SIZE:
+                try:
+                    buf = <WriteBuffer>next(bind_data)
+                except StopIteration:
+                    if packet.len() > 0:
+                        buffers.append(packet)
+                    if len(buffers) > 0:
+                        self._writelines(buffers)
+                        raise StopIteration(True)
+                    else:
+                        raise StopIteration(False)
+                except Exception as ex:
+                    raise StopIteration(ex)
+                buf = self._build_bind_message(portal_name, stmt_name, buf)
+                packet.write_buffer(buf)
+                buf = self._build_execute_message(portal_name, 0)
+                packet.write_buffer(buf)
+            buffers.append(packet)
+        self._writelines(buffers)
+
+    cdef _execute_many_done(self, bint data_sent):
+        if data_sent:
+            self._write(SYNC_MESSAGE)
+        else:
             self._push_result()
+
+    cdef _execute_many_fail(self, object error):
+        cdef WriteBuffer buf
+
+        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.is_in_transaction():
+            self._execute_many_done(True)
         else:
-            self._send_bind_message(portal_name, stmt_name, buf, 0)
+            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)
 
     cdef _execute(self, str portal_name, int32_t limit):
         cdef WriteBuffer buf
@@ -834,10 +868,7 @@ cdef class CoreProtocol:
 
         self.result = []
 
-        buf = WriteBuffer.new_message(b'E')
-        buf.write_str(portal_name, self.encoding)  # name of the portal
-        buf.write_int32(limit)  # number of rows to return; 0 - all
-        buf.end_message()
+        buf = self._build_execute_message(portal_name, limit)
 
         buf.write_bytes(SYNC_MESSAGE)
 
@@ -920,6 +951,9 @@ cdef class CoreProtocol:
     cdef _write(self, buf):
         raise NotImplementedError
 
+    cdef _writelines(self, list buffers):
+        raise NotImplementedError
+
     cdef _decode_row(self, const char* buf, ssize_t buf_len):
         pass
 

asyncpg/protocol/protocol.pyx

@@ -122,11 +122,6 @@ cdef class BaseProtocol(CoreProtocol):
     def get_settings(self):
         return self.settings
 
-    def is_in_transaction(self):
-        # PQTRANS_INTRANS = idle, within transaction block
-        # PQTRANS_INERROR = idle, within failed transaction
-        return self.xact_status in (PQTRANS_INTRANS, PQTRANS_INERROR)
-
     cdef inline resume_reading(self):
         if not self.is_reading:
             self.is_reading = True
@@ -216,15 +211,27 @@ cdef class BaseProtocol(CoreProtocol):
 
         waiter = self._new_waiter(timeout)
         try:
-            self._bind_execute_many(
-                portal_name,
-                state.name,
-                arg_bufs)  # network op
-
+            self._execute_many_init()
             self.last_query = state.query
             self.statement = state
             self.return_extra = False
             self.queries_count += 1
+
+            data_sent = False
+            while True:
+                self._execute_many_writelines(
+                    portal_name,
+                    state.name,
+                    arg_bufs)  # network op
+                data_sent = True
+                await self.writing_allowed.wait()
+        except StopIteration as ex:
+            if ex.value is True:
+                self._execute_many_done(True)  # network op
+            elif ex.value is False:
+                self._execute_many_done(data_sent)  # network op
+            else:
+                self._execute_many_fail(ex.value)  # network op
         except Exception as ex:
             waiter.set_exception(ex)
             self._coreproto_error()
@@ -880,6 +887,9 @@ cdef class BaseProtocol(CoreProtocol):
     cdef _write(self, buf):
         self.transport.write(memoryview(buf))
 
+    cdef _writelines(self, list buffers):
+        self.transport.writelines(buffers)
+
     # asyncio callbacks:
 
     def data_received(self, data):