From: Daniele Varrazzo Date: Sun, 2 Jan 2022 19:37:14 +0000 (+0100) Subject: Add returning parameter on executemany() X-Git-Tag: pool-3.1~49^2 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7aa500a548e9051c7f8d199fd6a7b9165f6817c4;p=thirdparty%2Fpsycopg.git Add returning parameter on executemany() --- diff --git a/docs/api/cursors.rst b/docs/api/cursors.rst index ea42d4dba..947e4f6df 100644 --- 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 ccee73bf6..76e857a61 100644 --- 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 cd3f04915..2c54d8bc9 100644 --- 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 1488b71c8..8ef1e9af9 100644 --- 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 293ff6ce4..fae1d70cc 100644 --- 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 da3f44290..fc526464a 100644 --- 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,)])