Fix KeyError in insert_all when pk column is missing from a single record

When a pk= column was named that isn't present in the data (and so isn't
created on the table), inserting exactly one record raised a KeyError while
inserting any other number of records did not. The single-record branch that
populates last_pk read the column straight out of the inserted row.

Check that the named pk column(s) actually exist on the row first and fall
back to the rowid otherwise, matching the behaviour for multiple records.

Closes #732

Author: ATOM00blue

sqlite_utils/db.py

@@ -3608,12 +3608,20 @@ def insert_all(
                 if (hash_id or 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:
-                        self.last_pk = row[hash_id]
-                    elif isinstance(pk, str):
-                        self.last_pk = row[pk]
+                    pk_cols = (
+                        [hash_id]
+                        if hash_id
+                        else ([pk] if isinstance(pk, str) else list(pk))
+                    )
+                    if all(col in row for col in pk_cols):
+                        if hash_id or isinstance(pk, str):
+                            self.last_pk = row[pk_cols[0]]
+                        else:
+                            self.last_pk = tuple(row[col] for col in pk_cols)
                     else:
-                        self.last_pk = tuple(row[p] for p in pk)
+                        # Named pk column(s) are not present in the table - fall
+                        # back to the rowid, matching the multi-row behaviour
+                        self.last_pk = self.last_rowid
                 else:
                     self.last_pk = self.last_rowid
             else:

tests/test_create.py

@@ -670,6 +670,22 @@ def test_insert_all_with_extra_columns_in_later_chunks(fresh_db):
     ]
 
 
+@pytest.mark.parametrize("num_rows", (0, 1, 2, 3, 10))
+def test_insert_all_pk_not_in_records(fresh_db, num_rows):
+    # https://github.com/simonw/sqlite-utils/issues/732
+    # Naming a pk= column that is absent from the records should behave the
+    # same regardless of how many rows are inserted - previously a single row
+    # raised a KeyError while other row counts did not.
+    fresh_db.conn.execute("CREATE TABLE t (a TEXT, b INT, PRIMARY KEY (a, b))")
+    rows = [{"a": "x{}".format(i), "b": i} for i in range(num_rows)]
+    table = fresh_db.table("t")
+    table.insert_all(rows, pk="not_a_column", alter=True)
+    assert table.count == num_rows
+    if num_rows == 1:
+        # Falls back to the rowid since the named pk column does not exist
+        assert table.last_pk == table.last_rowid
+
+
 def test_bulk_insert_more_than_999_values(fresh_db):
     "Inserting 100 items with 11 columns should work"
     fresh_db["big"].insert_all(