Fixing RTT in executemany

Author: fantix

asyncpg/protocol/coreproto.pxd

@@ -75,11 +75,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

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):
@@ -256,22 +251,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
@@ -799,27 +779,42 @@ cdef class CoreProtocol:
  cdef _bind_execute_many(self, str portal_name, str stmt_name,
  object bind_data):
 
- cdef WriteBuffer buf
+ cdef:
+ WriteBuffer packet
+ WriteBuffer buf
 
  self._ensure_connected()
  self._set_state(PROTOCOL_BIND_EXECUTE_MANY)
 
+ packet = WriteBuffer.new()
+
  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
- self._push_result()
- else:
- self._send_bind_message(portal_name, stmt_name, buf, 0)
+ while True:
+ try:
+ buf = <WriteBuffer>next(bind_data)
+ except StopIteration:
+ if packet.len() > 0:
+ packet.write_bytes(SYNC_MESSAGE)
+ self.transport.write(memoryview(packet))
+ else:
+ self._push_result()
+ break
+ except Exception as e:
+ self.result_type = RESULT_FAILED
+ self.result = e
+ self._push_result()
+ break
+ else:
+ buf = self._build_bind_message(portal_name, stmt_name, buf)
+ packet.write_buffer(buf)
+
+ buf = WriteBuffer.new_message(b'E')
+ buf.write_str(portal_name, self.encoding) # name of the portal
+ buf.write_int32(0) # number of rows to return; 0 - all
+ buf.end_message()
+ packet.write_buffer(buf)
 
  cdef _execute(self, str portal_name, int32_t limit):
  cdef WriteBuffer buf

tests/test_execute.py

@@ -151,3 +151,25 @@ async def test_execute_many_2(self):
  ''', good_data)
  finally:
  await self.con.execute('DROP TABLE exmany')
+
+ async def test_execute_many_atomic(self):
+ from asyncpg.exceptions import UniqueViolationError
+
+ await self.con.execute('CREATE TEMP TABLE exmany '
+ '(a text, b int PRIMARY KEY)')
+
+ try:
+ with self.assertRaises(UniqueViolationError):
+ await self.con.executemany('''
+ INSERT INTO exmany VALUES($1, $2)
+ ''', [
+ ('a', 1), ('b', 2), ('c', 2), ('d', 4)
+ ])
+
+ result = await self.con.fetch('''
+ SELECT * FROM exmany
+ ''')
+
+ self.assertEqual(result, [])
+ finally:
+ await self.con.execute('DROP TABLE exmany')