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

This commit is contained in:
Simon Willison 2025-11-23 10:57:08 -08:00
commit 7399ae48a1
2 changed files with 45 additions and 10 deletions

View file

@ -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()

View file

@ -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