From a29ba8cf34360e07c0d17876a04a51640f3e98ab Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 21 Sep 2020 19:40:38 -0700
Subject: [PATCH] .transform() can now alter NOT NULL

---
 sqlite_utils/db.py      | 19 ++++++++++++++++++-
 tests/test_transform.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index cfe8c01..05d5da7 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -735,9 +735,9 @@ class Table(Queryable):
             rename=rename,
             drop=None,
             pk=pk,
+            not_null=not_null,
             # foreign_keys=foreign_keys,
             # column_order=column_order,
-            # not_null=not_null,
             # defaults=defaults,
             # hash_id=hash_id,
             # extracts=extracts,
@@ -753,6 +753,7 @@ class Table(Queryable):
         rename=None,
         drop=None,
         pk=None,
+        not_null=None,
         tmp_suffix=None,
     ):
         columns = columns or {}
@@ -779,11 +780,27 @@ class Table(Queryable):
                 pk = self.pks[0]
             else:
                 pk = self.pks
+
+        # not_null may be a set or dict, need to convert to a set
+        create_table_not_null = {c.name for c in self.columns if c.notnull}
+        if isinstance(not_null, dict):
+            # Remove any columns with a value of False
+            for key, value in not_null.items():
+                # Column may have been renamed
+                key = rename.get(key) or key
+                if value is False and key in create_table_not_null:
+                    create_table_not_null.remove(key)
+                else:
+                    create_table_not_null.add(key)
+        else:
+            create_table_not_null.update(rename.get(k) or k for k in not_null)
+
         sqls.append(
             self.db.create_table_sql(
                 new_table_name,
                 dict(new_column_pairs),
                 pk=pk,
+                not_null=create_table_not_null,
                 # foreign_keys=foreign_keys,
                 # column_order=column_order,
                 # not_null=not_null,
diff --git a/tests/test_transform.py b/tests/test_transform.py
index 067a95b..e7a885e 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -98,3 +98,34 @@ def test_transform_sql_rowid_to_id(fresh_db):
         dogs.schema
         == 'CREATE TABLE "dogs" (\n   [id] INTEGER PRIMARY KEY,\n   [name] TEXT,\n   [age] TEXT\n)'
     )
+
+
+def test_transform_not_null(fresh_db):
+    dogs = fresh_db["dogs"]
+    dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id")
+    dogs.transform(not_null={"name"})
+    assert (
+        dogs.schema
+        == 'CREATE TABLE "dogs" (\n   [id] INTEGER PRIMARY KEY,\n   [name] TEXT NOT NULL,\n   [age] TEXT\n)'
+    )
+
+
+def test_transform_remove_a_not_null(fresh_db):
+    dogs = fresh_db["dogs"]
+    dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, not_null={"age"}, pk="id")
+    dogs.transform(not_null={"name": True, "age": False})
+    assert (
+        dogs.schema
+        == 'CREATE TABLE "dogs" (\n   [id] INTEGER PRIMARY KEY,\n   [name] TEXT NOT NULL,\n   [age] TEXT\n)'
+    )
+
+
+@pytest.mark.parametrize("not_null", [{"age"}, {"age": True}])
+def test_transform_add_not_null_with_rename(fresh_db, not_null):
+    dogs = fresh_db["dogs"]
+    dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id")
+    dogs.transform(not_null=not_null, rename={"age": "dog_age"})
+    assert (
+        dogs.schema
+        == 'CREATE TABLE "dogs" (\n   [id] INTEGER PRIMARY KEY,\n   [name] TEXT,\n   [dog_age] TEXT NOT NULL\n)'
+    )