Fix rowid pk and last_rowid regressions in insert/upsert

sqlite_utils/db.py

@@ -53,6 +53,11 @@
 
 SQLITE_MAX_VARS = 999
 
+# Names that refer to a rowid table's implicit integer primary key. These are
+# valid primary key targets even though they are not listed among a table's
+# columns. See https://www.sqlite.org/lang_createtable.html#rowid
+ROWID_ALIASES = frozenset({"rowid", "_rowid_", "oid"})
+
 _quote_fts_re = re.compile(r'\s+|(".*?")')
 
 _virtual_table_using_re = re.compile(
@@ -4343,10 +4348,14 @@ def insert_all(
         if pk and not hash_id and self.exists():
             pk_cols = [pk] if isinstance(pk, str) else list(pk)
             existing_columns = self.columns_dict
+            # rowid and its aliases are valid primary keys for a rowid table
+            # even though they are not listed among the table's columns
+            rowid_aliases = ROWID_ALIASES if self.use_rowid else frozenset()
             missing_pk_cols = [
                 col
                 for col in pk_cols
-                if resolve_casing(col, existing_columns) not in existing_columns
+                if col.lower() not in rowid_aliases
+                and resolve_casing(col, existing_columns) not in existing_columns
             ]
             if missing_pk_cols:
                 invalid_pk_error = InvalidColumns(
@@ -4512,40 +4521,72 @@ def insert_all(
             if not upsert and result is not None:
                 ignored_insert = ignore and result.rowcount == 0
                 if ignored_insert:
+                    # The row was not inserted because it conflicts with an
+                    # existing row. Point last_pk / last_rowid at that existing
+                    # row when we can identify it from the record's primary key
+                    # values, rather than leaving them stale or unset.
                     if list_mode:
-                        first_record_list = cast(Sequence[Any], first_record)
-                        if hash_id:
-                            pass
-                        elif isinstance(pk, str):
-                            pk_index = column_names.index(
-                                resolve_casing(pk, column_names)
-                            )
-                            self.last_pk = first_record_list[pk_index]
-                        elif pk:
-                            self.last_pk = tuple(
-                                first_record_list[
-                                    column_names.index(resolve_casing(p, column_names))
-                                ]
-                                for p in pk
-                            )
+                        first_record_dict = dict(
+                            zip(column_names, cast(Sequence[Any], first_record))
+                        )
                     else:
                         first_record_dict = cast(Dict[str, Any], first_record)
-                        if hash_id:
-                            self.last_pk = hash_record(
-                                first_record_dict, hash_id_columns
-                            )
-                        elif isinstance(pk, str):
-                            self.last_pk = first_record_dict[
-                                resolve_casing(pk, first_record_dict)
+                    if hash_id:
+                        self.last_pk = hash_record(first_record_dict, hash_id_columns)
+                    elif isinstance(pk, str):
+                        self.last_pk = first_record_dict[
+                            resolve_casing(pk, first_record_dict)
+                        ]
+                    elif pk:
+                        self.last_pk = tuple(
+                            first_record_dict[resolve_casing(p, first_record_dict)]
+                            for p in pk
+                        )
+                    # Locate the existing conflicting row using its primary key
+                    # columns so we can report its rowid (and pk if not already
+                    # known). Falls back to leaving them unset if the conflict
+                    # cannot be resolved to a pk lookup (e.g. a UNIQUE column).
+                    key_cols: Optional[List[str]] = None
+                    if isinstance(pk, str):
+                        key_cols = [pk]
+                    elif pk:
+                        key_cols = list(pk)
+                    elif not hash_id and not self.use_rowid:
+                        key_cols = self.pks
+                    if key_cols:
+                        try:
+                            key_values = [
+                                first_record_dict[resolve_casing(c, first_record_dict)]
+                                for c in key_cols
                             ]
-                        elif pk:
-                            self.last_pk = tuple(
-                                first_record_dict[resolve_casing(p, first_record_dict)]
-                                for p in pk
+                        except KeyError:
+                            key_values = None
+                        if key_values is not None:
+                            where = " and ".join(
+                                "{} = ?".format(quote_identifier(c)) for c in key_cols
                             )
+                            existing = self.db.execute(
+                                "select rowid from {} where {} limit 1".format(
+                                    quote_identifier(self.name), where
+                                ),
+                                key_values,
+                            ).fetchone()
+                            if existing is not None:
+                                self.last_rowid = existing[0]
+                                # On a primary key conflict the record's pk
+                                # values identify the existing row
+                                if self.last_pk is None:
+                                    self.last_pk = (
+                                        key_values[0]
+                                        if len(key_cols) == 1
+                                        else tuple(key_values)
+                                    )
                 else:
                     self.last_rowid = result.lastrowid
-                    if (hash_id or pk) and self.last_rowid:
+                    # A rowid-alias pk resolves directly to the rowid, so there
+                    # is no separate pk column to look up
+                    rowid_pk = isinstance(pk, str) and pk.lower() in ROWID_ALIASES
+                    if (hash_id or (pk and not rowid_pk)) and self.last_rowid:
                         # Set self.last_pk to the pk(s) for that rowid
                         row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
                         if hash_id:

tests/test_create.py

@@ -990,6 +990,96 @@ def test_insert_ignore(fresh_db):
     assert rows == [{"id": 1, "bar": 2}]
 
 
+def test_insert_ignore_reports_existing_row(fresh_db):
+    # An ignored insert (row already exists) should point last_rowid and
+    # last_pk at the existing conflicting row - see the Datasette insert API
+    fresh_db["docs"].insert({"id": 1, "title": "Exists"}, pk="id")
+    # Insert a conflicting row with ignore=True and no explicit pk=
+    table = fresh_db["docs"].insert({"id": 1, "title": "One"}, ignore=True)
+    assert table.last_rowid == 1
+    assert table.last_pk == 1
+    assert list(fresh_db["docs"].rows_where("rowid = ?", [table.last_rowid])) == [
+        {"id": 1, "title": "Exists"}
+    ]
+
+
+@pytest.mark.parametrize("rowid_alias", ("rowid", "_rowid_", "oid"))
+@pytest.mark.parametrize("method", ("upsert", "insert_replace", "insert_ignore"))
+def test_pk_rowid_alias_on_rowid_table(fresh_db, rowid_alias, method):
+    # rowid and its aliases are valid primary keys for a rowid table even
+    # though they are not listed among the table's columns - see the Datasette
+    # upsert API against tables without an explicit primary key
+    fresh_db["t"].insert({"title": "Hello"})
+    assert fresh_db["t"].pks == ["rowid"]
+    record = {rowid_alias: 1, "title": "Updated"}
+    if method == "upsert":
+        table = fresh_db["t"].upsert(record, pk=rowid_alias)
+    elif method == "insert_replace":
+        table = fresh_db["t"].insert(record, pk=rowid_alias, replace=True)
+    else:
+        table = fresh_db["t"].insert(record, pk=rowid_alias, ignore=True)
+    assert table.last_pk == 1
+    expected_title = "Hello" if method == "insert_ignore" else "Updated"
+    assert list(fresh_db["t"].rows) == [{"title": expected_title}]
+
+
+def test_insert_ignore_reports_existing_row_compound_pk(fresh_db):
+    # Compound primary key variant of the ignored-insert lookup
+    fresh_db["t"].insert_all([{"a": 1, "b": 2, "note": "first"}], pk=("a", "b"))
+    table = fresh_db["t"].insert(
+        {"a": 1, "b": 2, "note": "second"}, pk=("a", "b"), ignore=True
+    )
+    assert table.last_pk == (1, 2)
+    assert list(fresh_db["t"].rows_where("rowid = ?", [table.last_rowid])) == [
+        {"a": 1, "b": 2, "note": "first"}
+    ]
+
+
+def test_insert_ignore_reports_existing_row_list_mode(fresh_db):
+    # List-based iteration variant of the ignored-insert lookup
+    fresh_db["t"].insert_all([["id", "title"], [1, "first"]], pk="id")
+    table = fresh_db["t"].insert_all(
+        [["id", "title"], [1, "second"]], pk="id", ignore=True
+    )
+    assert table.last_pk == 1
+    assert table.last_rowid == 1
+    assert list(fresh_db["t"].rows) == [{"id": 1, "title": "first"}]
+
+
+def test_insert_ignore_hash_id_reports_pk(fresh_db):
+    # With hash_id the pk is the computed hash; the original record has no id
+    # column to look up so last_rowid is left unset
+    first = fresh_db["dogs"].insert({"name": "Cleo"}, hash_id="id")
+    table = fresh_db["dogs"].insert({"name": "Cleo"}, hash_id="id", ignore=True)
+    assert table.last_pk == first.last_pk
+    assert table.last_rowid is None
+    assert fresh_db["dogs"].count == 1
+
+
+def test_insert_ignore_unresolvable_conflict_leaves_pk_unset(fresh_db):
+    # When the conflict cannot be resolved to a primary key lookup, last_pk and
+    # last_rowid are left unset rather than reporting a misleading value
+
+    # rowid table with a UNIQUE column and no primary key: no pk to look up
+    fresh_db["u"].db.execute("create table u (title text unique)")
+    fresh_db["u"].insert({"title": "x"})
+    table = fresh_db["u"].insert({"title": "x"}, ignore=True)
+    assert table.last_pk is None
+    assert table.last_rowid is None
+    assert fresh_db["u"].count == 1
+
+    # Conflict on a UNIQUE column other than the primary key: the pk value from
+    # the record does not match the existing row, so the lookup finds nothing
+    fresh_db["docs"].db.execute(
+        "create table docs (id integer primary key, email text unique)"
+    )
+    fresh_db["docs"].insert({"id": 1, "email": "a"}, pk="id")
+    table = fresh_db["docs"].insert({"id": 2, "email": "a"}, ignore=True)
+    assert table.last_pk is None
+    assert table.last_rowid is None
+    assert fresh_db["docs"].count == 1
+
+
 def test_insert_ignore_with_pk_after_other_table_insert(fresh_db):
     # https://github.com/simonw/sqlite-utils/issues/554
     user = {"id": "abc", "name": "david"}