Add returning parameter on executemany()

diff --git a/docs/api/cursors.rst b/docs/api/cursors.rst
index ea42d4dbae052dce94283a49f18d0fe6834207c5..947e4f6df8f6802b0c5716e661a86908d7f7a99d 100644 (file)
--- a/docs/api/cursors.rst
+++ b/docs/api/cursors.rst
@@ -81,16 +81,29 @@ The `!Cursor` class
         :type query: `!str`, `!bytes`, or `sql.Composable`
         :param params_seq: The parameters to pass to the query
         :type params_seq: Sequence of Sequences or Mappings
+        :param returning: If `!false`, query results won't be available to fetch
+        :type returning: `!bool`
 
         This is more efficient than performing separate queries, but in case of
         several :sql:`INSERT` (and with some SQL creativity for massive
         :sql:`UPDATE` too) you may consider using `copy()`.
 
+        If the queries return data (e.g. when executing an :sql:`INSERT ...
+        RETURNING` or a :sql:`SELECT` with a side-effect), the result will be
+        available in the cursor's state, using `fetchone()` and similar
+        methods; results after the first will be available using `nextset()`.
+        In case this makes use of an unacceptable amount of memory you can use
+        `!returning=False` to disable returning the results. This is not
+        necessary if the queries return no data (e.g. on :sql:`INSERT` without
+        returning).
+
         See :ref:`query-parameters` for all the details about executing
         queries.
 
-        Results from each query, if any, can be walked through by calling
-        `nextset()`.
+        .. versionchanged:: 3.1
+
+            results are now available in the cursor state, added `!returning`
+            parameter to disable it.
 
     .. automethod:: copy
 

diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py
index ccee73bf64736a7d17656c4b127fe9d07a4bf51e..76e857a61e71e4b8bf7c5455b7f7a273b105a44b 100644 (file)
--- a/psycopg/psycopg/cursor.py
+++ b/psycopg/psycopg/cursor.py
@@ -207,7 +207,7 @@ class BaseCursor(Generic[ConnectionType, Row]):
             yield from self._conn._exec_command(cmd)
 
     def _executemany_gen(
-        self, query: Query, params_seq: Iterable[Params]
+        self, query: Query, params_seq: Iterable[Params], returning: bool
     ) -> PQGen[None]:
         """Generator implementing `Cursor.executemany()`."""
         yield from self._start_query(query)
@@ -223,7 +223,8 @@ class BaseCursor(Generic[ConnectionType, Row]):
 
             results = yield from self._maybe_prepare_gen(pgq, prepare=True)
             self._check_results(results)
-            self._results.extend(results)
+            if returning and results[0].status == ExecStatus.TUPLES_OK:
+                self._results.extend(results)
 
             for res in results:
                 nrows += res.command_tuples or 0
@@ -577,13 +578,21 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
             raise ex.with_traceback(None)
         return self
 
-    def executemany(self, query: Query, params_seq: Iterable[Params]) -> None:
+    def executemany(
+        self,
+        query: Query,
+        params_seq: Iterable[Params],
+        *,
+        returning: bool = True,
+    ) -> None:
         """
         Execute the same command with a sequence of input data.
         """
         try:
             with self._conn.lock:
-                self._conn.wait(self._executemany_gen(query, params_seq))
+                self._conn.wait(
+                    self._executemany_gen(query, params_seq, returning)
+                )
         except e.Error as ex:
             raise ex.with_traceback(None)
 

diff --git a/psycopg/psycopg/cursor_async.py b/psycopg/psycopg/cursor_async.py
index cd3f04915dd5f3611df57425a8ece225f9e0d8b3..2c54d8bc9d17e46b1f6f082e1003ef718eeb17ea 100644 (file)
--- a/psycopg/psycopg/cursor_async.py
+++ b/psycopg/psycopg/cursor_async.py
@@ -82,11 +82,17 @@ class AsyncCursor(BaseCursor["AsyncConnection[Any]", Row]):
         return self
 
     async def executemany(
-        self, query: Query, params_seq: Iterable[Params]
+        self,
+        query: Query,
+        params_seq: Iterable[Params],
+        *,
+        returning: bool = True,
     ) -> None:
         try:
             async with self._conn.lock:
-                await self._conn.wait(self._executemany_gen(query, params_seq))
+                await self._conn.wait(
+                    self._executemany_gen(query, params_seq, returning)
+                )
         except e.Error as ex:
             raise ex.with_traceback(None)
 

diff --git a/psycopg/psycopg/server_cursor.py b/psycopg/psycopg/server_cursor.py
index 1488b71c8f6b3a54d8d02fcf5c349d74a9c35cb1..8ef1e9af9f6706b1df1af7955d837e4dcd507efa 100644 (file)
--- a/psycopg/psycopg/server_cursor.py
+++ b/psycopg/psycopg/server_cursor.py
@@ -266,7 +266,13 @@ class ServerCursor(Cursor[Row]):
 
         return self
 
-    def executemany(self, query: Query, params_seq: Iterable[Params]) -> None:
+    def executemany(
+        self,
+        query: Query,
+        params_seq: Iterable[Params],
+        *,
+        returning: bool = True,
+    ) -> None:
         """Method not implemented for server-side cursors."""
         raise e.NotSupportedError(
             "executemany not supported on server-side cursors"
@@ -387,7 +393,11 @@ class AsyncServerCursor(AsyncCursor[Row]):
         return self
 
     async def executemany(
-        self, query: Query, params_seq: Iterable[Params]
+        self,
+        query: Query,
+        params_seq: Iterable[Params],
+        *,
+        returning: bool = True,
     ) -> None:
         raise e.NotSupportedError(
             "executemany not supported on server-side cursors"

diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index 293ff6ce44d6fc36bd979988a6d9fe5e735c85f8..fae1d70cc1727976af9e7f1e5b714256f96dfb5a 100644 (file)
--- a/tests/test_cursor.py
+++ b/tests/test_cursor.py
@@ -295,6 +295,31 @@ def test_executemany_returning(conn, execmany):
     assert cur.nextset() is None
 
 
+def test_executemany_returning_discard(conn, execmany):
+    cur = conn.cursor()
+    cur.executemany(
+        "insert into execmany(num, data) values (%s, %s) returning num",
+        [(10, "hello"), (20, "world")],
+        returning=False,
+    )
+    assert cur.rowcount == 2
+    with pytest.raises(psycopg.ProgrammingError):
+        cur.fetchone()
+    assert cur.nextset() is None
+
+
+def test_executemany_no_result(conn, execmany):
+    cur = conn.cursor()
+    cur.executemany(
+        "insert into execmany(num, data) values (%s, %s)",
+        [(10, "hello"), (20, "world")],
+    )
+    assert cur.rowcount == 2
+    with pytest.raises(psycopg.ProgrammingError):
+        cur.fetchone()
+    assert cur.nextset() is None
+
+
 def test_executemany_rowcount_no_hit(conn, execmany):
     cur = conn.cursor()
     cur.executemany("delete from execmany where id = %s", [(-1,), (-2,)])

diff --git a/tests/test_cursor_async.py b/tests/test_cursor_async.py
index da3f442903f9daf129d7a39e5c29f06331c43db8..fc526464ac9b41ec2691d59658e8a31e14552fc8 100644 (file)
--- a/tests/test_cursor_async.py
+++ b/tests/test_cursor_async.py
@@ -287,6 +287,31 @@ async def test_executemany_returning(aconn, execmany):
     assert cur.nextset() is None
 
 
+async def test_executemany_returning_discard(aconn, execmany):
+    cur = aconn.cursor()
+    await cur.executemany(
+        "insert into execmany(num, data) values (%s, %s) returning num",
+        [(10, "hello"), (20, "world")],
+        returning=False,
+    )
+    assert cur.rowcount == 2
+    with pytest.raises(psycopg.ProgrammingError):
+        await cur.fetchone()
+    assert cur.nextset() is None
+
+
+async def test_executemany_no_result(aconn, execmany):
+    cur = aconn.cursor()
+    await cur.executemany(
+        "insert into execmany(num, data) values (%s, %s)",
+        [(10, "hello"), (20, "world")],
+    )
+    assert cur.rowcount == 2
+    with pytest.raises(psycopg.ProgrammingError):
+        await cur.fetchone()
+    assert cur.nextset() is None
+
+
 async def test_executemany_rowcount_no_hit(aconn, execmany):
     cur = aconn.cursor()
     await cur.executemany("delete from execmany where id = %s", [(-1,), (-2,)])