From 2376c452a56b0c3e75e7ca698273434e32945304 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 8 May 2023 12:24:10 -0700
Subject: [PATCH] upsert_all() now works with not_null - refs #538

---
 sqlite_utils/db.py   | 22 +++++++++++++++++-----
 tests/test_upsert.py | 10 ++++++++++
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index c06e6a0..44d59da 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -2741,6 +2741,7 @@ class Table(Queryable):
         hash_id_columns,
         upsert,
         pk,
+        not_null,
         conversions,
         num_records_processed,
         replace,
@@ -2778,14 +2779,20 @@ class Table(Queryable):
                 pks = pk
             self.last_pk = None
             for record_values in values:
-                # TODO: make more efficient:
                 record = dict(zip(all_columns, record_values))
-                sql = "INSERT OR IGNORE INTO [{table}]({pks}) VALUES({pk_placeholders});".format(
+                placeholders = list(pks)
+                # Need to populate not-null columns too, or INSERT OR IGNORE ignores
+                # them since it ignores the resulting integrity errors
+                if not_null:
+                    placeholders.extend(not_null)
+                sql = "INSERT OR IGNORE INTO [{table}]({cols}) VALUES({placeholders});".format(
                     table=self.name,
-                    pks=", ".join(["[{}]".format(p) for p in pks]),
-                    pk_placeholders=", ".join(["?" for p in pks]),
+                    cols=", ".join(["[{}]".format(p) for p in placeholders]),
+                    placeholders=", ".join(["?" for p in placeholders]),
+                )
+                queries_and_params.append(
+                    (sql, [record[col] for col in pks] + ["" for _ in (not_null or [])])
                 )
-                queries_and_params.append((sql, [record[col] for col in pks]))
                 # UPDATE [book] SET [name] = 'Programming' WHERE [id] = 1001;
                 set_cols = [col for col in all_columns if col not in pks]
                 if set_cols:
@@ -2846,6 +2853,7 @@ class Table(Queryable):
         hash_id_columns,
         upsert,
         pk,
+        not_null,
         conversions,
         num_records_processed,
         replace,
@@ -2859,6 +2867,7 @@ class Table(Queryable):
             hash_id_columns,
             upsert,
             pk,
+            not_null,
             conversions,
             num_records_processed,
             replace,
@@ -2888,6 +2897,7 @@ class Table(Queryable):
                             hash_id_columns,
                             upsert,
                             pk,
+                            not_null,
                             conversions,
                             num_records_processed,
                             replace,
@@ -2903,6 +2913,7 @@ class Table(Queryable):
                             hash_id_columns,
                             upsert,
                             pk,
+                            not_null,
                             conversions,
                             num_records_processed,
                             replace,
@@ -3112,6 +3123,7 @@ class Table(Queryable):
                 hash_id_columns,
                 upsert,
                 pk,
+                not_null,
                 conversions,
                 num_records_processed,
                 replace,
diff --git a/tests/test_upsert.py b/tests/test_upsert.py
index cef8af0..b7267ea 100644
--- a/tests/test_upsert.py
+++ b/tests/test_upsert.py
@@ -28,6 +28,16 @@ def test_upsert_all_single_column(fresh_db):
     assert table.pks == ["name"]
 
 
+def test_upsert_all_not_null(fresh_db):
+    # https://github.com/simonw/sqlite-utils/issues/538
+    fresh_db["comments"].upsert_all(
+        [{"id": 1, "name": "Cleo"}],
+        pk="id",
+        not_null=["name"],
+    )
+    assert list(fresh_db["comments"].rows) == [{"id": 1, "name": "Cleo"}]
+
+
 def test_upsert_error_if_no_pk(fresh_db):
     table = fresh_db["table"]
     with pytest.raises(PrimaryKeyRequired):