Add fetchmany to execute many *and* return rows (#1175)

Co-authored-by: Elvis Pranskevichus <elvis@edgedb.com>

asyncpg/connection.py

@@ -756,6 +756,44 @@ async def fetchrow(
  return None
  return data[0]
 
+ async def fetchmany(
+ self, query, args, *, timeout: float=None, record_class=None
+ ):
+ """Run a query for each sequence of arguments in *args*
+ and return the results as a list of :class:`Record`.
+
+ :param query:
+ Query to execute.
+ :param args:
+ An iterable containing sequences of arguments for the query.
+ :param float timeout:
+ Optional timeout value in seconds.
+ :param type record_class:
+ If specified, the class to use for records returned by this method.
+ Must be a subclass of :class:`~asyncpg.Record`. If not specified,
+ a per-connection *record_class* is used.
+
+ :return list:
+ A list of :class:`~asyncpg.Record` instances. If specified, the
+ actual type of list elements would be *record_class*.
+
+ Example:
+
+ .. code-block:: pycon
+
+ >>> rows = await con.fetchmany('''
+ ... INSERT INTO mytab (a, b) VALUES ($1, $2) RETURNING a;
+ ... ''', [('x', 1), ('y', 2), ('z', 3)])
+ >>> rows
+ [<Record row=('x',)>, <Record row=('y',)>, <Record row=('z',)>]
+
+ .. versionadded:: 0.30.0
+ """
+ self._check_open()
+ return await self._executemany(
+ query, args, timeout, return_rows=True, record_class=record_class
+ )
+
  async def copy_from_table(self, table_name, *, output,
  columns=None, schema_name=None, timeout=None,
  format=None, oids=None, delimiter=None,
@@ -1896,17 +1934,27 @@ async def __execute(
  )
  return result, stmt
 
- async def _executemany(self, query, args, timeout):
+ async def _executemany(
+ self,
+ query,
+ args,
+ timeout,
+ return_rows=False,
+ record_class=None,
+ ):
  executor = lambda stmt, timeout: self._protocol.bind_execute_many(
  state=stmt,
  args=args,
  portal_name='',
  timeout=timeout,
+ return_rows=return_rows,
  )
  timeout = self._protocol._get_timeout(timeout)
  with self._stmt_exclusive_section:
  with self._time_and_log(query, args, timeout):
- result, _ = await self._do_execute(query, executor, timeout)
+ result, _ = await self._do_execute(
+ query, executor, timeout, record_class=record_class
+ )
  return result
 
  async def _do_execute(

asyncpg/pool.py

@@ -609,6 +609,22 @@ async def fetchrow(self, query, *args, timeout=None, record_class=None):
  record_class=record_class
  )
 
+ async def fetchmany(self, query, args, *, timeout=None, record_class=None):
+ """Run a query for each sequence of arguments in *args*
+ and return the results as a list of :class:`Record`.
+
+ Pool performs this operation using one of its connections. Other than
+ that, it behaves identically to
+ :meth:`Connection.fetchmany()
+ <asyncpg.connection.Connection.fetchmany>`.
+
+ .. versionadded:: 0.30.0
+ """
+ async with self.acquire() as con:
+ return await con.fetchmany(
+ query, args, timeout=timeout, record_class=record_class
+ )
+
  async def copy_from_table(
  self,
  table_name,

asyncpg/prepared_stmt.py

@@ -210,6 +210,27 @@ async def fetchrow(self, *args, timeout=None):
  return None
  return data[0]
 
+ @connresource.guarded
+ async def fetchmany(self, args, *, timeout=None):
+ """Execute the statement and return a list of :class:`Record` objects.
+
+ :param args: Query arguments.
+ :param float timeout: Optional timeout value in seconds.
+
+ :return: A list of :class:`Record` instances.
+
+ .. versionadded:: 0.30.0
+ """
+ return await self.__do_execute(
+ lambda protocol: protocol.bind_execute_many(
+ self._state,
+ args,
+ portal_name='',
+ timeout=timeout,
+ return_rows=True,
+ )
+ )
+
  @connresource.guarded
  async def executemany(self, args, *, timeout: float=None):
  """Execute the statement for each sequence of arguments in *args*.
@@ -222,7 +243,12 @@ async def executemany(self, args, *, timeout: float=None):
  """
  return await self.__do_execute(
  lambda protocol: protocol.bind_execute_many(
- self._state, args, '', timeout))
+ self._state,
+ args,
+ portal_name='',
+ timeout=timeout,
+ return_rows=False,
+ ))
 
  async def __do_execute(self, executor):
  protocol = self._connection._protocol

asyncpg/protocol/coreproto.pxd

@@ -171,7 +171,7 @@ cdef class CoreProtocol:
  cdef _bind_execute(self, str portal_name, str stmt_name,
  WriteBuffer bind_data, int32_t limit)
  cdef bint _bind_execute_many(self, str portal_name, str stmt_name,
- object bind_data)
+ object bind_data, bint return_rows)
  cdef bint _bind_execute_many_more(self, bint first=*)
  cdef _bind_execute_many_fail(self, object error, bint first=*)
  cdef _bind(self, str portal_name, str stmt_name,

asyncpg/protocol/coreproto.pyx

@@ -1020,12 +1020,12 @@ cdef class CoreProtocol:
  self._send_bind_message(portal_name, stmt_name, bind_data, limit)
 
  cdef bint _bind_execute_many(self, str portal_name, str stmt_name,
- object bind_data):
+ object bind_data, bint return_rows):
  self._ensure_connected()
  self._set_state(PROTOCOL_BIND_EXECUTE_MANY)
 
- self.result = None
- self._discard_data = True
+ self.result = [] if return_rows else None
+ self._discard_data = not return_rows
  self._execute_iter = bind_data
  self._execute_portal_name = portal_name
  self._execute_stmt_name = stmt_name

asyncpg/protocol/protocol.pyx

@@ -212,6 +212,7 @@ cdef class BaseProtocol(CoreProtocol):
  args,
  portal_name: str,
  timeout,
+ return_rows: bool,
  ):
  if self.cancel_waiter is not None:
  await self.cancel_waiter
@@ -237,7 +238,8 @@ cdef class BaseProtocol(CoreProtocol):
  more = self._bind_execute_many(
  portal_name,
  state.name,
- arg_bufs) # network op
+ arg_bufs,
+ return_rows) # network op
 
  self.last_query = state.query
  self.statement = state

tests/test_execute.py

@@ -139,6 +139,45 @@ async def test_executemany_basic(self):
  ('a', 1), ('b', 2), ('c', 3), ('d', 4)
  ])
 
+ async def test_executemany_returning(self):
+ result = await self.con.fetchmany('''
+ INSERT INTO exmany VALUES($1, $2) RETURNING a, b
+ ''', [
+ ('a', 1), ('b', 2), ('c', 3), ('d', 4)
+ ])
+ self.assertEqual(result, [
+ ('a', 1), ('b', 2), ('c', 3), ('d', 4)
+ ])
+ result = await self.con.fetch('''
+ SELECT * FROM exmany
+ ''')
+ self.assertEqual(result, [
+ ('a', 1), ('b', 2), ('c', 3), ('d', 4)
+ ])
+
+ # Empty set
+ await self.con.fetchmany('''
+ INSERT INTO exmany VALUES($1, $2) RETURNING a, b
+ ''', ())
+ result = await self.con.fetch('''
+ SELECT * FROM exmany
+ ''')
+ self.assertEqual(result, [
+ ('a', 1), ('b', 2), ('c', 3), ('d', 4)
+ ])
+
+ # Without "RETURNING"
+ result = await self.con.fetchmany('''
+ INSERT INTO exmany VALUES($1, $2)
+ ''', [('e', 5), ('f', 6)])
+ self.assertEqual(result, [])
+ result = await self.con.fetch('''
+ SELECT * FROM exmany
+ ''')
+ self.assertEqual(result, [
+ ('a', 1), ('b', 2), ('c', 3), ('d', 4), ('e', 5), ('f', 6)
+ ])
+
  async def test_executemany_bad_input(self):
  with self.assertRaisesRegex(
  exceptions.DataError,

tests/test_prepare.py

@@ -611,3 +611,17 @@ async def test_prepare_explicitly_named(self):
  'prepared statement "foobar" already exists',
  ):
  await self.con.prepare('select 1', name='foobar')
+
+ async def test_prepare_fetchmany(self):
+ tr = self.con.transaction()
+ await tr.start()
+ try:
+ await self.con.execute('CREATE TABLE fetchmany (a int, b text)')
+
+ stmt = await self.con.prepare(
+ 'INSERT INTO fetchmany (a, b) VALUES ($1, $2) RETURNING a, b'
+ )
+ result = await stmt.fetchmany([(1, 'a'), (2, 'b'), (3, 'c')])
+ self.assertEqual(result, [(1, 'a'), (2, 'b'), (3, 'c')])
+ finally:
+ await tr.rollback()