diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index da06238..e97f71d 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -961,7 +961,7 @@ class Table: else: raise self.last_rowid = result.lastrowid - self.last_pk = None + self.last_pk = self.last_rowid # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened if (hash_id or pk) and self.last_rowid: row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0] diff --git a/tests/test_update.py b/tests/test_update.py new file mode 100644 index 0000000..0b99333 --- /dev/null +++ b/tests/test_update.py @@ -0,0 +1,24 @@ +import pytest + + +def test_update_rowid_table(fresh_db): + table = fresh_db["table"] + rowid = table.insert({"foo": "bar"}).last_pk + table.update(rowid, {"foo": "baz"}) + assert [{"foo": "baz"}] == list(table.rows) + + +def test_update_pk_table(fresh_db): + table = fresh_db["table"] + pk = table.insert({"foo": "bar", "id": 5}, pk="id").last_pk + assert 5 == pk + table.update(pk, {"foo": "baz"}) + assert [{"id": 5, "foo": "baz"}] == list(table.rows) + + +def test_update_compound_pk_table(fresh_db): + table = fresh_db["table"] + pk = table.insert({"id1": 5, "id2": 3, "v": 1}, pk=("id1", "id2")).last_pk + assert (5, 3) == pk + table.update(pk, {"v": 2}) + assert [{"id1": 5, "id2": 3, "v": 2}] == list(table.rows)