From bfd74a35bb055a2dd28da8eabd149e4e03f30fed Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 24 Jun 2026 10:43:53 -0700 Subject: [PATCH] upsert() detects existing compound primary keys Closes #629 --- sqlite_utils/db.py | 5 +++++ tests/test_upsert.py | 47 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 38e95c7..ae99322 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -3554,6 +3554,11 @@ class Table(Queryable): if hash_id_columns and hash_id is None: hash_id = "id" + if upsert and not pk and not hash_id and self.exists(): + existing_pks = [column.name for column in self.columns if column.is_pk] + if existing_pks: + pk = existing_pks[0] if len(existing_pks) == 1 else tuple(existing_pks) + if upsert and (not pk and not hash_id): raise PrimaryKeyRequired("upsert() requires a pk") diff --git a/tests/test_upsert.py b/tests/test_upsert.py index de0e362..09783bb 100644 --- a/tests/test_upsert.py +++ b/tests/test_upsert.py @@ -49,6 +49,53 @@ def test_upsert_error_if_no_pk(fresh_db): table.upsert({"id": 1, "name": "Cleo"}) +def test_upsert_error_if_existing_table_has_no_pk(fresh_db): + table = fresh_db.create_table("table", {"id": int, "name": str}) + with pytest.raises(PrimaryKeyRequired): + table.upsert({"id": 1, "name": "Cleo"}) + + +@pytest.mark.parametrize("use_old_upsert", (False, True)) +def test_upsert_uses_compound_pk_from_existing_table(use_old_upsert): + # https://github.com/simonw/sqlite-utils/issues/629 + db = Database(memory=True, use_old_upsert=use_old_upsert) + db.execute(""" + create table summary ( + Source text, + Object text, + Category text, + Count integer, + primary key (Source, Object, Category) + ) + """) + table = db["summary"] + table.upsert( + { + "Source": "Client A", + "Object": "Accounts", + "Category": "All", + "Count": 3, + } + ) + assert table.last_pk == ("Client A", "Accounts", "All") + table.upsert( + { + "Source": "Client A", + "Object": "Accounts", + "Category": "All", + "Count": 4, + } + ) + assert list(table.rows) == [ + { + "Source": "Client A", + "Object": "Accounts", + "Category": "All", + "Count": 4, + } + ] + + def test_upsert_with_hash_id(fresh_db): table = fresh_db["table"] table.upsert({"foo": "bar"}, hash_id="pk")