Issue #16864: Cursor.lastrowid now supports REPLACE statement

berkerpeksag · berkerpeksag · commit e0b70cd8a986 · 2016-06-14T15:25:36.000+03:00
Initial patch by Alex LordThorsen.
diff --git a/Doc/library/sqlite3.rst b/Doc/library/sqlite3.rst
@@ -629,9 +629,16 @@ Cursor Objects
    .. attribute:: lastrowid
 
       This read-only attribute provides the rowid of the last modified row. It is
-      only set if you issued an ``INSERT`` statement using the :meth:`execute`
-      method. For operations other than ``INSERT`` or when :meth:`executemany` is
-      called, :attr:`lastrowid` is set to :const:`None`.
+      only set if you issued an ``INSERT`` or a ``REPLACE`` statement using the
+      :meth:`execute` method.  For operations other than ``INSERT`` or
+      ``REPLACE`` or when :meth:`executemany` is called, :attr:`lastrowid` is
+      set to :const:`None`.
+
+      If the ``INSERT`` or ``REPLACE`` statement failed to insert the previous
+      successful rowid is returned.
+
+      .. versionchanged:: 3.6
+         Added support for the ``REPLACE`` statement.
 
    .. attribute:: description
 
diff --git a/Doc/whatsnew/3.6.rst b/Doc/whatsnew/3.6.rst
@@ -339,6 +339,13 @@ you may now specify file paths on top of directories (e.g. zip files).
 (Contributed by Wolfgang Langner in :issue:`26587`).
 
 
+sqlite3
+-------
+
+* :attr:`sqlite3.Cursor.lastrowid` now supports the ``REPLACE`` statement.
+  (Contributed by Alex LordThorsen in :issue:`16864`.)
+
+
 socketserver
 ------------
 
diff --git a/Lib/sqlite3/test/dbapi.py b/Lib/sqlite3/test/dbapi.py
@@ -188,7 +188,10 @@ class CursorTests(unittest.TestCase):
     def setUp(self):
         self.cx = sqlite.connect(":memory:")
         self.cu = self.cx.cursor()
-        self.cu.execute("create table test(id integer primary key, name text, income number)")
+        self.cu.execute(
+            "create table test(id integer primary key, name text, "
+            "income number, unique_test text unique)"
+        )
         self.cu.execute("insert into test(name) values (?)", ("foo",))
 
     def tearDown(self):
@@ -462,6 +465,44 @@ class Foo: pass
         with self.assertRaises(TypeError):
             cur = sqlite.Cursor(foo)
 
+    def CheckLastRowIDOnReplace(self):
+        """
+        INSERT OR REPLACE and REPLACE INTO should produce the same behavior.
+        """
+        sql = '{} INTO test(id, unique_test) VALUES (?, ?)'
+        for statement in ('INSERT OR REPLACE', 'REPLACE'):
+            with self.subTest(statement=statement):
+                self.cu.execute(sql.format(statement), (1, 'foo'))
+                self.assertEqual(self.cu.lastrowid, 1)
+
+    def CheckLastRowIDOnIgnore(self):
+        self.cu.execute(
+            "insert or ignore into test(unique_test) values (?)",
+            ('test',))
+        self.assertEqual(self.cu.lastrowid, 2)
+        self.cu.execute(
+            "insert or ignore into test(unique_test) values (?)",
+            ('test',))
+        self.assertEqual(self.cu.lastrowid, 2)
+
+    def CheckLastRowIDInsertOR(self):
+        results = []
+        for statement in ('FAIL', 'ABORT', 'ROLLBACK'):
+            sql = 'INSERT OR {} INTO test(unique_test) VALUES (?)'
+            with self.subTest(statement='INSERT OR {}'.format(statement)):
+                self.cu.execute(sql.format(statement), (statement,))
+                results.append((statement, self.cu.lastrowid))
+                with self.assertRaises(sqlite.IntegrityError):
+                    self.cu.execute(sql.format(statement), (statement,))
+                results.append((statement, self.cu.lastrowid))
+        expected = [
+            ('FAIL', 2), ('FAIL', 2),
+            ('ABORT', 3), ('ABORT', 3),
+            ('ROLLBACK', 4), ('ROLLBACK', 4),
+        ]
+        self.assertEqual(results, expected)
+
+
 @unittest.skipUnless(threading, 'This test requires threading.')
 class ThreadTests(unittest.TestCase):
     def setUp(self):
diff --git a/Misc/ACKS b/Misc/ACKS
@@ -895,6 +895,7 @@ Martin von Löwis
 Hugo Lopes Tavares
 Guillermo López-Anglada
 Anne Lord
+Alex LordThorsen
 Tom Loredo
 Justin Love
 Ned Jackson Lovely
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -10,6 +10,9 @@ What's New in Python 3.6.0 alpha 3
 Library
 -------
 
+- Issue #16864: sqlite3.Cursor.lastrowid now supports REPLACE statement.
+  Initial patch by Alex LordThorsen.
+
 - Issue #26386: Fixed ttk.TreeView selection operations with item id's
   containing spaces.
 
diff --git a/Modules/_sqlite/cursor.c b/Modules/_sqlite/cursor.c
@@ -698,7 +698,9 @@ PyObject* _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject*
         }
 
         Py_DECREF(self->lastrowid);
-        if (!multiple && statement_type == STATEMENT_INSERT) {
+        if (!multiple &&
+            /* REPLACE is an alias for INSERT OR REPLACE */
+            (statement_type == STATEMENT_INSERT || statement_type == STATEMENT_REPLACE)) {
             sqlite_int64 lastrowid;
             Py_BEGIN_ALLOW_THREADS
             lastrowid = sqlite3_last_insert_rowid(self->connection->db);
