Refs MagicStack#289, rewrite bind_execute_many()

Author: fantix

asyncpg/connection.py

@@ -252,14 +252,24 @@ async def executemany(self, command: str, args, *, timeout: float=None):
             ... ''', [(1, 2, 3), (4, 5, 6)])
 
         :param command: Command to execute.
-        :param args: An iterable containing sequences of arguments.
+        :param args: An (async) 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.7.0
 
         .. versionchanged:: 0.11.0
            `timeout` became a keyword-only parameter.
+
+        .. versionchanged:: 0.16.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.
+
+        .. versionchanged:: 0.16.0
+           ``args`` could be an async iterable.
         """
         self._check_open()
         return await self._executemany(command, args, timeout)

asyncpg/protocol/consts.pxi

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

asyncpg/protocol/coreproto.pxd

@@ -31,6 +31,7 @@ cdef enum ProtocolState:
     PROTOCOL_COPY_OUT_DONE = 20
     PROTOCOL_COPY_IN = 21
     PROTOCOL_COPY_IN_DATA = 22
+    PROTOCOL_BIND_EXECUTE_MANY_ROLLBACK = 23
 
 
 cdef enum AuthenticationMessage:
@@ -75,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
@@ -158,8 +154,10 @@ 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_send(self, object buffers)
+    cdef _execute_many_rollback(self, object error)
+    cdef _execute_many_end(self, object sync)
     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

@@ -25,11 +25,6 @@ cdef class CoreProtocol:
 
         self._skip_discard = False
 
-        # executemany support data
-        self._execute_iter = None
-        self._execute_portal_name = None
-        self._execute_stmt_name = None
-
         self._reset_result()
 
     cdef _write(self, buf):
@@ -70,7 +65,8 @@ cdef class CoreProtocol:
                 elif state == PROTOCOL_BIND_EXECUTE:
                     self._process__bind_execute(mtype)
 
-                elif state == PROTOCOL_BIND_EXECUTE_MANY:
+                elif (state == PROTOCOL_BIND_EXECUTE_MANY or
+                          state == PROTOCOL_BIND_EXECUTE_MANY_ROLLBACK):
                     self._process__bind_execute_many(mtype)
 
                 elif state == PROTOCOL_EXECUTE:
@@ -247,7 +243,9 @@ cdef class CoreProtocol:
 
         elif mtype == b'E':
             # ErrorResponse
-            self._parse_msg_error_response(True)
+            # Ignoring expected rollback error
+            self._parse_msg_error_response(
+                self.state != PROTOCOL_BIND_EXECUTE_MANY_ROLLBACK)
 
         elif mtype == b'2':
             # BindComplete
@@ -256,22 +254,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
@@ -670,6 +653,10 @@ cdef class CoreProtocol:
                     new_state == PROTOCOL_COPY_IN_DATA):
                 self.state = new_state
 
+            elif (self.state == PROTOCOL_BIND_EXECUTE_MANY and
+                    new_state == PROTOCOL_BIND_EXECUTE_MANY_ROLLBACK):
+                self.state = new_state
+
             elif self.state == PROTOCOL_FAILED:
                 raise apg_exc.InternalClientError(
                     'cannot switch to state {}; '
@@ -811,30 +798,55 @@ 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:
+    cdef _execute_many_send(self, object buffers):
+        self._ensure_connected()
+        self.transport.writelines(buffers)
+
+    cdef _execute_many_rollback(self, object error):
+        cdef:
+            WriteBuffer packet
+            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
+
+    cdef _execute_many_end(self, object sync):
+        if sync is True:
+            self._write_sync_message()
+        elif sync is False:
             self._push_result()
-        except Exception as e:
+        else:
             self.result_type = RESULT_FAILED
-            self.result = e
+            self.result = sync
             self._push_result()
-        else:
-            self._send_bind_message(portal_name, stmt_name, buf, 0)
 
     cdef _execute(self, str portal_name, int32_t limit):
         cdef WriteBuffer buf

asyncpg/protocol/protocol.pyx

@@ -208,6 +208,9 @@ cdef class BaseProtocol(CoreProtocol):
     @cython.iterable_coroutine
     async def bind_execute_many(self, PreparedStatementState state, args,
                                 str portal_name, timeout):
+        cdef:
+            WriteBuffer buf
+            WriteBuffer packet
 
         if self.cancel_waiter is not None:
             await self.cancel_waiter
@@ -217,29 +220,128 @@ cdef class BaseProtocol(CoreProtocol):
 
         self._check_state()
         timeout = self._get_timeout_impl(timeout)
+        timer = Timer(timeout)
 
-        # Make sure the argument sequence is encoded lazily with
-        # this generator expression to keep the memory pressure under
-        # control.
-        data_gen = (state._encode_bind_msg(b) for b in args)
-        arg_bufs = iter(data_gen)
+        try:
+            aiter = args.__aiter__
+        except AttributeError:
+            iterator = args.__iter__()
+            async_iterator = False
+        else:
+            iterator = aiter()
+            async_iterator = True
+
+        waiter = self._new_waiter(timer.get_remaining_budget())
+
+        self._execute_many_init()
+        self.last_query = state.query
+        self.statement = state
+        self.return_extra = False
+        self.queries_count += 1
+
+        more = True  # for marking end of input
+        should_commit = False  # for marking empty input
+        rollback_on_error = False  # for input fails before send to DB
 
-        waiter = self._new_waiter(timeout)
         try:
-            self._bind_execute_many(
-                portal_name,
-                state.name,
-                arg_bufs)  # network op
+            # Send all input data until end of iterable reached
+            while more:
 
-            self.last_query = state.query
-            self.statement = state
-            self.return_extra = False
-            self.queries_count += 1
-        except Exception as ex:
-            waiter.set_exception(ex)
-            self._coreproto_error()
-        finally:
-            return await waiter
+                # Fill up 4 buffers to take advantage of uvloop `writelines()`
+                buffers = []
+                while more and len(buffers) < _EXECUTE_MANY_BUF_NUM:
+
+                    # A new buffer
+                    packet = WriteBuffer.new()
+
+                    try:
+                        # Fill up this buffer up to 32KB or a bit more
+                        while more and packet.len() < _EXECUTE_MANY_BUF_SIZE:
+
+                            # Grab the next item from input
+                            if async_iterator:
+                                with timer:
+                                    args_item = await asyncio.wait_for(
+                                        iterator.__anext__(),
+                                        timeout=timer.get_remaining_budget(),
+                                        loop=self.loop)
+
+                                # If server issues an error in parallel, we
+                                # abort with SYNC now
+                                if self.result_type == RESULT_FAILED:
+                                    packet.write_bytes(SYNC_MESSAGE)
+                                    more = False
+                                    break
+                            else:
+                                args_item = iterator.__next__()
+
+                            # Encode the arguments
+                            args_buf = state._encode_bind_msg(args_item)
+
+                            # And write into the buffer on protocol
+                            buf = self._build_bind_message(
+                                portal_name, state.name, args_buf)
+                            packet.write_buffer(buf)
+                            buf = self._build_execute_message(portal_name, 0)
+                            packet.write_buffer(buf)
+
+                            # Mark that we should send SYNC in the end
+                            should_commit = True
+
+                    except (builtins.StopIteration,
+                            builtins.StopAsyncIteration):
+
+                        if should_commit:
+                            # End of input reached, append SYNC in current
+                            # buffer and mark the end to escape loops
+                            packet.write_bytes(SYNC_MESSAGE)
+                            more = False
+
+                        else:
+                            # Break loops and handle later to unify results
+                            raise
+
+                    # One buffer ready, 3/2/1/no more to go
+                    buffers.append(packet)
+
+                # All 4 buffers are ready now, send them once wire has capacity
+                with timer:
+                    await self.writing_allowed.wait()
+
+                if self.result_type == RESULT_FAILED:
+                    # In any case if server returns an error, we should stop
+                    # sending more data, and abort with a SYNC message
+                    self._execute_many_end(True)  # network op
+                    more = False
+
+                else:
+                    # This uses `loop.writelines()` for its possible efficiency
+                    self._execute_many_send(buffers)  # network op
+
+                    # If input fails, we need to rollback from now on
+                    rollback_on_error = True
+
+        except (builtins.StopIteration, builtins.StopAsyncIteration):
+            # Input is empty, just reset the state
+            self._execute_many_end(False)
+
+        except Exception as e:
+            # It fails for:
+            # 1. input iteration failure;
+            # 2. timeout;
+            # 3. other protocol code error.
+
+            if rollback_on_error:
+                # We sent some data previously without SYNC, for atomicity we
+                # should rollback, wait for ReadyForQuery then re-raise
+                self._execute_many_rollback(e)  # network op
+
+            else:
+                # If we never sent anything to DB, it is safe to reset the
+                # state and re-raise
+                self._execute_many_end(e)
+
+        return await waiter
 
     @cython.iterable_coroutine
     async def bind(self, PreparedStatementState state, args,