From: Denis Laxalde Date: Sat, 7 Jan 2023 15:01:36 +0000 (+0100) Subject: fix: set rowcount to the first result in executemany(..., returning=True) X-Git-Tag: pool-3.2.0~136^2 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=refs%2Fpull%2F479%2Fhead;p=thirdparty%2Fpsycopg.git fix: set rowcount to the first result in executemany(..., returning=True) Previously, _rowcount was unconditionally set to the overall number of rows of all queries in executemany() and then reset only upon the first call to nextset(). In the returning=True case, this lead the rowcount attribute to be wrong for the first result (i.e. it didn't match the number of rows that would be returned by .fetchall(), as can be seen in updated tests). Now we only set _rowcount to the cumulated number of rows of executed queries *if* executemany() is not returning (so the value is still useful, e.g., in to check the number of INSERTed rows): >>> cur.executemany("INSERT INTO t(r) VALUES (%s)", [(1,), (2,)]) >>> cur.rowcount 2 # number of inserted rows >>> cur.nextset() >>> cur.executemany("INSERT INTO t(r) VALUES (%s) RETURNING r", [(1,), (2,)], returning=True) >>> cur.rowcount 1 # number of rows in the first result set >>> cur.fetchall() [(1,)] >>> cur.nextset() True >>> cur.rowcount 1 >>> cur.fetchall() [(2,)] >>> cur.nextset() Besides, the code for processing results from executemany() in _executemany_gen_no_pipeline() is now similar to that of _set_results_from_pipeline(). --- diff --git a/docs/api/cursors.rst b/docs/api/cursors.rst index 9c5b47813..78c336bb0 100644 --- a/docs/api/cursors.rst +++ b/docs/api/cursors.rst @@ -105,6 +105,11 @@ The `!Cursor` class methods. Each input parameter will produce a separate result set: use `nextset()` to read the results of the queries after the first one. + The value of `rowcount` is set to the cumulated number of rows + affected by queries; except when using `!returning=True`, in which + case it is set to the number of rows in the current result set (i.e. + the first one, until `nextset()` gets called). + See :ref:`query-parameters` for all the details about executing queries. @@ -239,6 +244,10 @@ The `!Cursor` class a successful command, such as ``CREATE TABLE`` or ``UPDATE 42``. .. autoattribute:: rowcount + + From `executemany()`, unless called with `!returning=True`, this is + the cumulated number of rows affected by executed commands. + .. autoattribute:: rownumber .. attribute:: _query diff --git a/docs/news.rst b/docs/news.rst index 69c626b1f..471d3cd3a 100644 --- a/docs/news.rst +++ b/docs/news.rst @@ -15,6 +15,8 @@ Psycopg 3.1.8 (unreleased) - Don't pollute server logs when types looked for by `TypeInfo.fetch()` are not found (:ticket:`#473`). +- Set `Cursor.rowcount` to the number of rows of each result set from + `~Cursor.executemany()` when called with `!returning=True` (:ticket:`#479`). Current release --------------- diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py index 08cdcb751..7c32f29e4 100644 --- a/psycopg/psycopg/cursor.py +++ b/psycopg/psycopg/cursor.py @@ -219,7 +219,8 @@ class BaseCursor(Generic[ConnectionType, Row]): assert pipeline yield from self._start_query(query) - self._rowcount = 0 + if not returning: + self._rowcount = 0 assert self._execmany_returning is None self._execmany_returning = returning @@ -251,8 +252,9 @@ class BaseCursor(Generic[ConnectionType, Row]): Generator implementing `Cursor.executemany()` with pipelines not available. """ yield from self._start_query(query) + if not returning: + self._rowcount = 0 first = True - nrows = 0 for params in params_seq: if first: pgq = self._convert_query(query, params) @@ -266,17 +268,15 @@ class BaseCursor(Generic[ConnectionType, Row]): self._check_results(results) if returning: self._results.extend(results) - - for res in results: - nrows += res.command_tuples or 0 + else: + # In non-returning case, set rowcount to the cumulated number + # of rows of executed queries. + for res in results: + self._rowcount += res.command_tuples or 0 if self._results: self._select_current_result(0) - # Override rowcount for the first result. Calls to nextset() will change - # it to the value of that result only, but we hope nobody will notice. - # You haven't read this comment. - self._rowcount = nrows self._last_query = query for cmd in self._conn._prepared.get_maintenance_commands(): @@ -545,14 +545,11 @@ class BaseCursor(Generic[ConnectionType, Row]): self._results.extend(results) if first_batch: self._select_current_result(0) - self._rowcount = 0 - - # Override rowcount for the first result. Calls to nextset() will - # change it to the value of that result only, but we hope nobody - # will notice. - # You haven't read this comment. - for res in results: - self._rowcount += res.command_tuples or 0 + else: + # In non-returning case, set rowcount to the cumulated number of + # rows of executed queries. + for res in results: + self._rowcount += res.command_tuples or 0 def _send_prepare(self, name: bytes, query: PostgresQuery) -> None: if self._conn._pipeline: diff --git a/tests/test_client_cursor.py b/tests/test_client_cursor.py index b35560474..e091e9cf1 100644 --- a/tests/test_client_cursor.py +++ b/tests/test_client_cursor.py @@ -326,9 +326,10 @@ def test_executemany_returning(conn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.fetchone() == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert cur.fetchone() == (20,) assert cur.nextset() is None @@ -352,12 +353,13 @@ def test_executemany_no_result(conn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") with pytest.raises(psycopg.ProgrammingError): cur.fetchone() pgresult = cur.pgresult assert cur.nextset() + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") assert pgresult is not cur.pgresult assert cur.nextset() is None diff --git a/tests/test_client_cursor_async.py b/tests/test_client_cursor_async.py index 0cf8ec649..25f4810c0 100644 --- a/tests/test_client_cursor_async.py +++ b/tests/test_client_cursor_async.py @@ -316,9 +316,10 @@ async def test_executemany_returning(aconn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert (await cur.fetchone()) == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert (await cur.fetchone()) == (20,) assert cur.nextset() is None @@ -342,12 +343,13 @@ async def test_executemany_no_result(aconn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") with pytest.raises(psycopg.ProgrammingError): await cur.fetchone() pgresult = cur.pgresult assert cur.nextset() + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") assert pgresult is not cur.pgresult assert cur.nextset() is None diff --git a/tests/test_cursor.py b/tests/test_cursor.py index a667f4fb3..a39ed6755 100644 --- a/tests/test_cursor.py +++ b/tests/test_cursor.py @@ -308,9 +308,10 @@ def test_executemany_returning(conn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.fetchone() == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert cur.fetchone() == (20,) assert cur.nextset() is None @@ -334,12 +335,13 @@ def test_executemany_no_result(conn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") with pytest.raises(psycopg.ProgrammingError): cur.fetchone() pgresult = cur.pgresult assert cur.nextset() + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") assert pgresult is not cur.pgresult assert cur.nextset() is None diff --git a/tests/test_cursor_async.py b/tests/test_cursor_async.py index ac3fdeb2c..fdc3d898f 100644 --- a/tests/test_cursor_async.py +++ b/tests/test_cursor_async.py @@ -295,9 +295,10 @@ async def test_executemany_returning(aconn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert (await cur.fetchone()) == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert (await cur.fetchone()) == (20,) assert cur.nextset() is None @@ -321,12 +322,13 @@ async def test_executemany_no_result(aconn, execmany): [(10, "hello"), (20, "world")], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") with pytest.raises(psycopg.ProgrammingError): await cur.fetchone() pgresult = cur.pgresult assert cur.nextset() + assert cur.rowcount == 1 assert cur.statusmessage.startswith("INSERT") assert pgresult is not cur.pgresult assert cur.nextset() is None diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 56fe59888..3ef4014fd 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -326,9 +326,10 @@ def test_executemany(conn): [(10,), (20,)], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert cur.fetchone() == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert cur.fetchone() == (20,) assert cur.nextset() is None diff --git a/tests/test_pipeline_async.py b/tests/test_pipeline_async.py index 2e743cfa2..1dc611085 100644 --- a/tests/test_pipeline_async.py +++ b/tests/test_pipeline_async.py @@ -327,9 +327,10 @@ async def test_executemany(aconn): [(10,), (20,)], returning=True, ) - assert cur.rowcount == 2 + assert cur.rowcount == 1 assert (await cur.fetchone()) == (10,) assert cur.nextset() + assert cur.rowcount == 1 assert (await cur.fetchone()) == (20,) assert cur.nextset() is None