Fix for bug where self.last_pk didn't work for single list rows

Author: simonw

sqlite_utils/db.py

@@ -3529,17 +3529,29 @@ def insert_all(
                     self.last_pk = self.last_rowid
             else:
                 # For an upsert use first_record from earlier
-                # Note: This code path assumes dict mode; list mode upserts
-                # with single records don't populate last_pk correctly yet
-                first_record_dict = cast(Dict[str, Any], first_record)
-                if hash_id:
-                    self.last_pk = hash_record(first_record_dict, hash_id_columns)
+                if list_mode:
+                    # In list mode, look up pk value by column index
+                    first_record_list = cast(Sequence[Any], first_record)
+                    if hash_id:
+                        # hash_id not supported in list mode for last_pk
+                        pass
+                    elif isinstance(pk, str):
+                        pk_index = column_names.index(pk)
+                        self.last_pk = first_record_list[pk_index]
+                    else:
+                        self.last_pk = tuple(
+                            first_record_list[column_names.index(p)] for p in pk
+                        )
                 else:
-                    self.last_pk = (
-                        first_record_dict[pk]
-                        if isinstance(pk, str)
-                        else tuple(first_record_dict[p] for p in pk)
-                    )
+                    first_record_dict = cast(Dict[str, Any], first_record)
+                    if hash_id:
+                        self.last_pk = hash_record(first_record_dict, hash_id_columns)
+                    else:
+                        self.last_pk = (
+                            first_record_dict[pk]
+                            if isinstance(pk, str)
+                            else tuple(first_record_dict[p] for p in pk)
+                        )
 
         if analyze:
             self.analyze()

tests/test_list_mode.py

@@ -263,3 +263,26 @@ def data_generator():
     assert rows[0] == {"id": 1, "name": "Alice", "age": 30, "city": "NYC"}
     assert rows[1] == {"id": 2, "name": "Bob", "age": None, "city": None}
     assert rows[2] == {"id": 3, "name": "Charlie", "age": 35, "city": None}
+
+
+def test_list_mode_single_record_upsert_last_pk():
+    """Test that last_pk is populated correctly for single-record upserts in list mode"""
+    db = Database(memory=True)
+
+    # Create table first
+    db["data"].insert({"id": 1, "name": "Alice", "value": 100}, pk="id")
+
+    # Now upsert a single record using list mode
+    def upsert_data():
+        yield ["id", "name", "value"]
+        yield [1, "Alice", 150]  # Update existing
+
+    table = db["data"]
+    table.upsert_all(upsert_data(), pk="id")
+
+    # Verify the data was updated
+    rows = list(db["data"].rows)
+    assert rows == [{"id": 1, "name": "Alice", "value": 150}]
+
+    # Verify last_pk is populated correctly
+    assert table.last_pk == 1