[Python-checkins] [3.10] gh-93421: Update sqlite3 cursor.rowcount only after SQLITE_DONE (GH-93526) (GH-93599)

[erlend-aasland]

https://github.com/python/cpython/commit/9cc0afcb8782c22575fa91756c54029c7e774dfb
commit: 9cc0afcb8782c22575fa91756c54029c7e774dfb
branch: 3.10
author: Erlend Egeberg Aasland <erlend.aasland at protonmail.com>
committer: erlend-aasland <erlend.aasland at protonmail.com>
date: 2022-06-08T13:31:14+02:00
summary:

[3.10] gh-93421: Update sqlite3 cursor.rowcount only after SQLITE_DONE (GH-93526) (GH-93599)

(cherry picked from commit 875de61)

files:
A Misc/NEWS.d/next/Library/2022-06-05-22-22-42.gh-issue-93421.43UO_8.rst
M Lib/sqlite3/test/dbapi.py
M Modules/_sqlite/cursor.c

diff --git a/Lib/sqlite3/test/dbapi.py b/Lib/sqlite3/test/dbapi.py
index e332184a7d107..a764c82b0d30f 100644
--- a/Lib/sqlite3/test/dbapi.py
+++ b/Lib/sqlite3/test/dbapi.py
@@ -381,6 +381,14 @@ def test_rowcount_executemany(self):
         self.cu.executemany("insert into test(name) values (?)", [(1,), (2,), (3,)])
         self.assertEqual(self.cu.rowcount, 3)
 
+    @unittest.skipIf(sqlite.sqlite_version_info < (3, 35, 0),
+                     "Requires SQLite 3.35.0 or newer")
+    def test_rowcount_update_returning(self):
+        # gh-93421: rowcount is updated correctly for UPDATE...RETURNING queries
+        self.cu.execute("update test set name='bar' where name='foo' returning 1")
+        self.assertEqual(self.cu.fetchone()[0], 1)
+        self.assertEqual(self.cu.rowcount, 1)
+
     def test_total_changes(self):
         self.cu.execute("insert into test(name) values ('foo')")
         self.cu.execute("insert into test(name) values ('foo')")
diff --git a/Misc/NEWS.d/next/Library/2022-06-05-22-22-42.gh-issue-93421.43UO_8.rst b/Misc/NEWS.d/next/Library/2022-06-05-22-22-42.gh-issue-93421.43UO_8.rst
new file mode 100644
index 0000000000000..9e1d6554e0ab2
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-06-05-22-22-42.gh-issue-93421.43UO_8.rst
@@ -0,0 +1,3 @@
+Update :data:`sqlite3.Cursor.rowcount` when a DML statement has run to
+completion. This fixes the row count for SQL queries like
+``UPDATE ... RETURNING``. Patch by Erlend E. Aasland.
diff --git a/Modules/_sqlite/cursor.c b/Modules/_sqlite/cursor.c
index ac80c285fe995..ca42af4fed403 100644
--- a/Modules/_sqlite/cursor.c
+++ b/Modules/_sqlite/cursor.c
@@ -492,10 +492,9 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
         pysqlite_statement_reset(self->statement);
     }
 
-    /* reset description and rowcount */
+    /* reset description */
     Py_INCREF(Py_None);
     Py_SETREF(self->description, Py_None);
-    self->rowcount = 0L;
 
     func_args = PyTuple_New(1);
     if (!func_args) {
@@ -527,6 +526,7 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
 
     pysqlite_statement_reset(self->statement);
     pysqlite_statement_mark_dirty(self->statement);
+    self->rowcount = self->statement->is_dml ? 0L : -1L;
 
     /* We start a transaction implicitly before a DML statement.
        SELECT is the only exception. See #9924. */
@@ -604,12 +604,6 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             }
         }
 
-        if (self->statement->is_dml) {
-            self->rowcount += (long)sqlite3_changes(self->connection->db);
-        } else {
-            self->rowcount= -1L;
-        }
-
         if (!multiple) {
             Py_BEGIN_ALLOW_THREADS
             lastrowid = sqlite3_last_insert_rowid(self->connection->db);
@@ -630,11 +624,17 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             if (self->next_row == NULL)
                 goto error;
         } else if (rc == SQLITE_DONE && !multiple) {
+            if (self->statement->is_dml) {
+                self->rowcount = (long)sqlite3_changes(self->connection->db);
+            }
             pysqlite_statement_reset(self->statement);
             Py_CLEAR(self->statement);
         }
 
         if (multiple) {
+            if (self->statement->is_dml && rc == SQLITE_DONE) {
+                self->rowcount += (long)sqlite3_changes(self->connection->db);
+            }
             pysqlite_statement_reset(self->statement);
         }
         Py_XDECREF(parameters);
@@ -824,7 +824,12 @@ pysqlite_cursor_iternext(pysqlite_Cursor *self)
         if (PyErr_Occurred()) {
             goto error;
         }
-        if (rc != SQLITE_DONE && rc != SQLITE_ROW) {
+        if (rc == SQLITE_DONE) {
+            if (self->statement->is_dml) {
+                self->rowcount = (long)sqlite3_changes(self->connection->db);
+            }
+        }
+        else if (rc != SQLITE_ROW) {
             _pysqlite_seterror(self->connection->db, NULL);
             goto error;
         }