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 9c5b4781391da7c890409d750179bdaa29787d4f..78c336bb00b265f288c218cd28d55c625a4b0e52 100644 (file)
--- 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 69c626b1f8ffba8ecdf6e936fda2255c6a38a16f..471d3cd3a7a41e0c51cebc208d30bee5f124ca2f 100644 (file)
--- 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 08cdcb7510c5798895d6ce5cabb818b400ead739..7c32f29e4f67087f02fdafdc2ef8bed440657732 100644 (file)
--- 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 b3556047410114404eb59679342edbd7258b65c1..e091e9cf1025f62f2473ae55583b9422000929fd 100644 (file)
--- 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 0cf8ec6494a09bd4526b61c28af10c15c3c1c4f2..25f4810c0fe5f211a5d11d788d2e44e8473c378b 100644 (file)
--- 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 a667f4fb32d917e0a9537808e0fbc0acf9e3341d..a39ed6755e614a3ed9788b4505af07875a5817cb 100644 (file)
--- 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 ac3fdeb2cb9903ec36f183f042c73c381e6f14bb..fdc3d898f4d985f21f0d270b67c5c094e04aa134 100644 (file)
--- 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 56fe59888eab33c2023e750f875e92914c56f58d..3ef4014fd7e8bda1d76b67b37384653687d45611 100644 (file)
--- 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 2e743cfa2476ee705522c48757c6f624f5ec99f1..1dc6110858db0bb183427b23d0a624a514977277 100644 (file)
--- 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