From 489eda92bc3b528c90b76ac90a3f9b78c8ea02a3 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 5 Jan 2020 09:20:11 -0800
Subject: [PATCH] .upsert() and upsert_all() require pk=, closes #73

---
 sqlite_utils/db.py   |  9 ++++++---
 tests/test_upsert.py | 12 ++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 7516028..0a81a60 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -90,6 +90,10 @@ class NotFoundError(Exception):
     pass
 
 
+class PrimaryKeyRequired(Exception):
+    pass
+
+
 class Database:
     def __init__(self, filename_or_conn=None, memory=False):
         assert (filename_or_conn is not None and not memory) or (
@@ -957,6 +961,8 @@ class Table(Queryable):
         data
         """
         pk = self.value_or_default("pk", pk)
+        if upsert and not pk:
+            raise PrimaryKeyRequired("upsert() requires a pk")
         foreign_keys = self.value_or_default("foreign_keys", foreign_keys)
         column_order = self.value_or_default("column_order", column_order)
         not_null = self.value_or_default("not_null", not_null)
@@ -1142,9 +1148,6 @@ class Table(Queryable):
         alter=DEFAULT,
         extracts=DEFAULT,
     ):
-        # Perform the following for each record:
-        # INSERT OR IGNORE INTO books(id) VALUES(1001);
-        # UPDATE books SET name = 'Programming' WHERE id = 1001;
         return self.insert_all(
             records,
             pk=pk,
diff --git a/tests/test_upsert.py b/tests/test_upsert.py
index edf7f29..4d53fa8 100644
--- a/tests/test_upsert.py
+++ b/tests/test_upsert.py
@@ -1,3 +1,7 @@
+from sqlite_utils.db import PrimaryKeyRequired
+import pytest
+
+
 def test_upsert(fresh_db):
     table = fresh_db["table"]
     table.insert({"id": 1, "name": "Cleo"}, pk="id")
@@ -16,6 +20,14 @@ def test_upsert_all(fresh_db):
     assert 2 == table.last_pk
 
 
+def test_upsert_error_if_no_pk(fresh_db):
+    table = fresh_db["table"]
+    with pytest.raises(PrimaryKeyRequired):
+        table.upsert_all([{"id": 1, "name": "Cleo"}])
+    with pytest.raises(PrimaryKeyRequired):
+        table.upsert({"id": 1, "name": "Cleo"})
+
+
 def test_upsert_compound_primary_key(fresh_db):
     table = fresh_db["table"]
     table.upsert_all(