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

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(