diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index a9c9eb5..dd83b4c 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -3529,17 +3529,29 @@ class Table(Queryable): 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() diff --git a/tests/test_list_mode.py b/tests/test_list_mode.py index 01f0a49..746c9c1 100644 --- a/tests/test_list_mode.py +++ b/tests/test_list_mode.py @@ -263,3 +263,26 @@ def test_tuple_mode_shorter_rows(): 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