From 32f8173a8fe830c224e39a0a514cd12e78de7028 Mon Sep 17 00:00:00 2001 From: Colin Dellow Date: Sat, 26 Nov 2022 11:11:06 -0500 Subject: [PATCH] upsert new rows with constraints --- sqlite_utils/db.py | 11 +++++++---- tests/test_upsert.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index e819d17..27dc52c 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -2779,13 +2779,16 @@ class Table(Queryable): self.last_pk = None for record_values in values: # TODO: make more efficient: + + # The initial insert statement includes all columns, so that upserts + # of new rows whose non-pk columns have constraints can succeed record = dict(zip(all_columns, record_values)) - sql = "INSERT OR IGNORE INTO [{table}]({pks}) VALUES({pk_placeholders});".format( + sql = "INSERT OR IGNORE INTO [{table}]({columns}) VALUES({column_placeholders});".format( table=self.name, - pks=", ".join(["[{}]".format(p) for p in pks]), - pk_placeholders=", ".join(["?" for p in pks]), + columns=", ".join(["[{}]".format(p) for p in all_columns]), + column_placeholders=", ".join(["?" for p in all_columns]), ) - queries_and_params.append((sql, [record[col] for col in pks])) + queries_and_params.append((sql, [record[col] for col in all_columns])) # UPDATE [book] SET [name] = 'Programming' WHERE [id] = 1001; set_cols = [col for col in all_columns if col not in pks] if set_cols: diff --git a/tests/test_upsert.py b/tests/test_upsert.py index cef8af0..7c80a4c 100644 --- a/tests/test_upsert.py +++ b/tests/test_upsert.py @@ -36,6 +36,20 @@ def test_upsert_error_if_no_pk(fresh_db): table.upsert({"id": 1, "name": "Cleo"}) +def test_upsert_with_constraints(fresh_db): + table = fresh_db.create_table( + "table_with_constraints", + { + "id": "text", + "name": "text", + }, + not_null=["name"], + ) + + table.upsert({"id": 1, "name": "Cleo"}, pk="id") + assert 1 == table.last_pk + + def test_upsert_with_hash_id(fresh_db): table = fresh_db["table"] table.upsert({"foo": "bar"}, hash_id="pk")