From 635c91475aa52e58b467797a95fec4554908f7dc Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 12 Apr 2020 20:22:32 -0700 Subject: [PATCH] Only set last_pk on singular .insert()/.update(), refs #98 --- sqlite_utils/db.py | 37 ++++++++++++++++++++++++++----------- tests/test_create.py | 7 ++++++- tests/test_upsert.py | 11 +++++++++-- 3 files changed, 41 insertions(+), 14 deletions(-) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 8d22e0c..d51ede3 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -468,6 +468,9 @@ class Queryable: class Table(Queryable): + last_rowid = None + last_pk = None + def __init__( self, db, @@ -987,6 +990,7 @@ class Table(Queryable): ), "Use either ignore=True or replace=True, not both" all_columns = None first = True + num_records_processed = 0 # We can only handle a max of 999 variables in a SQL insert, so # we need to adjust the batch_size down if we have too many cols records = iter(records) @@ -1000,8 +1004,11 @@ class Table(Queryable): num_columns <= SQLITE_MAX_VARS ), "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS) batch_size = max(1, min(batch_size, SQLITE_MAX_VARS // num_columns)) + self.last_rowid = None + self.last_pk = None for chunk in chunks(itertools.chain([first_record], records), batch_size): chunk = list(chunk) + num_records_processed += len(chunk) if first: if not self.exists(): # Use the first batch to derive the table names @@ -1046,6 +1053,7 @@ class Table(Queryable): pks = [pk] else: pks = pk + self.last_pk = None for record_values in values: # TODO: make more efficient: record = dict(zip(all_columns, record_values)) @@ -1073,6 +1081,12 @@ class Table(Queryable): + [record[pk] for pk in pks], ) ) + # We can populate .last_pk right here + if num_records_processed == 1: + self.last_pk = tuple(record[pk] for pk in pks) + if len(self.last_pk) == 1: + self.last_pk = self.last_pk[0] + else: or_what = "" if replace: @@ -1110,17 +1124,18 @@ class Table(Queryable): result = self.db.conn.execute(query, params) else: raise - self.last_rowid = result.lastrowid - 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] - if hash_id: - self.last_pk = row[hash_id] - elif isinstance(pk, str): - self.last_pk = row[pk] - else: - self.last_pk = tuple(row[p] for p in pk) + if num_records_processed == 1 and not upsert: + self.last_rowid = result.lastrowid + 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] + if hash_id: + self.last_pk = row[hash_id] + elif isinstance(pk, str): + self.last_pk = row[pk] + else: + self.last_pk = tuple(row[p] for p in pk) return self def upsert( diff --git a/tests/test_create.py b/tests/test_create.py index 905f49d..7edcd8d 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -135,7 +135,12 @@ def test_create_table_with_not_null(fresh_db): ), ) def test_create_table_from_example(fresh_db, example, expected_columns): - fresh_db["people"].insert(example) + people_table = fresh_db["people"] + assert None == people_table.last_rowid + assert None == people_table.last_pk + people_table.insert(example) + assert 1 == people_table.last_rowid + assert 1 == people_table.last_pk assert ["people"] == fresh_db.table_names() assert expected_columns == [ {"name": col.name, "type": col.type} for col in fresh_db["people"].columns diff --git a/tests/test_upsert.py b/tests/test_upsert.py index 584a496..100c6c1 100644 --- a/tests/test_upsert.py +++ b/tests/test_upsert.py @@ -7,6 +7,7 @@ def test_upsert(fresh_db): table.insert({"id": 1, "name": "Cleo"}, pk="id") table.upsert({"id": 1, "age": 5}, pk="id", alter=True) assert [{"id": 1, "name": "Cleo", "age": 5}] == list(table.rows) + assert 1 == table.last_pk def test_upsert_all(fresh_db): @@ -17,7 +18,7 @@ def test_upsert_all(fresh_db): {"id": 1, "name": "Cleo", "age": 5}, {"id": 2, "name": "Nixie", "age": 5}, ] == list(table.rows) - assert 2 == table.last_pk + assert table.last_pk is None def test_upsert_error_if_no_pk(fresh_db): @@ -34,6 +35,7 @@ def test_upsert_with_hash_id(fresh_db): assert [{"pk": "a5e744d0164540d33b1d7ea616c28f2fa97e754a", "foo": "bar"}] == list( table.rows ) + assert "a5e744d0164540d33b1d7ea616c28f2fa97e754a" == table.last_pk def test_upsert_compound_primary_key(fresh_db): @@ -45,8 +47,13 @@ def test_upsert_compound_primary_key(fresh_db): ], pk=("species", "id"), ) - table.upsert_all([{"species": "dog", "id": 1, "age": 5}], pk=("species", "id")) + assert None == table.last_pk + table.upsert({"species": "dog", "id": 1, "age": 5}, pk=("species", "id")) + assert ("dog", 1) == table.last_pk assert [ {"species": "dog", "id": 1, "name": "Cleo", "age": 5}, {"species": "cat", "id": 1, "name": "Catbag", "age": None}, ] == list(table.rows) + # .upsert_all() with a single item should set .last_pk + table.upsert_all([{"species": "cat", "id": 1, "age": 5}], pk=("species", "id")) + assert ("cat", 1) == table.last_pk