From 2fed87da6ea990d295672e4db2c8ae97b787913e Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 12 Jun 2019 18:35:02 -0700
Subject: [PATCH] --not-null-default and not_null_default=, refs #24

---
 docs/cli.rst         |   4 ++++
 docs/python-api.rst  |   6 ++++++
 sqlite_utils/cli.py  |  23 +++++++++++++++++++----
 sqlite_utils/db.py   |  15 +++++++++++++--
 tests/.DS_Store      | Bin 0 -> 6148 bytes
 tests/test_cli.py    |  21 +++++++++++++++++++++
 tests/test_create.py |  34 ++++++++++++++++++++++++----------
 7 files changed, 87 insertions(+), 16 deletions(-)
 create mode 100644 tests/.DS_Store

diff --git a/docs/cli.rst b/docs/cli.rst
index 4ef6584..27994ba 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -280,6 +280,10 @@ You can explicitly specify the column you wish to reference using ``--fk-col``::
 
     $ sqlite-utils add-column mydb.db dogs species_id --fk species --fk-col ref
 
+You can set a ``NOT NULL DEFAULT 'x'`` constraint on the new column using ``--not-null-default``::
+
+    $ sqlite-utils add-column mydb.db dogs friends_count integer --not-null-default 0
+
 .. _cli_add_column_alter:
 
 Adding columns automatically on insert/update
diff --git a/docs/python-api.rst b/docs/python-api.rst
index 331a84e..41f13d8 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -284,6 +284,12 @@ You can explicitly specify the column you wish to reference using ``fk_col``:
 
     db["dogs"].add_column("species_id", fk="species", fk_col="ref")
 
+You can set a ``NOT NULL DEFAULT 'x'`` constraint on the new column using ``not_null_default``:
+
+.. code-block:: python
+
+    db["dogs"].add_column("friends_count", int, not_null_default=0)
+
 .. _python_api_add_column_alter:
 
 Adding columns automatically on insert/update
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index daf4abc..4a32ddb 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -169,12 +169,27 @@ def optimize(path, no_vacuum):
     ),
     required=False,
 )
-@click.option("--fk", type=str, required=False)
-@click.option("--fk-col", type=str, required=False)
-def add_column(path, table, col_name, col_type, fk, fk_col):
+@click.option(
+    "--fk", type=str, required=False, help="Table to reference as a foreign key"
+)
+@click.option(
+    "--fk-col",
+    type=str,
+    required=False,
+    help="Referenced column on that foreign key table - if omitted will automatically use the primary key",
+)
+@click.option(
+    "--not-null-default",
+    type=str,
+    required=False,
+    help="Add NOT NULL DEFAULT 'TEXT' constraint",
+)
+def add_column(path, table, col_name, col_type, fk, fk_col, not_null_default):
     "Add a column to the specified table"
     db = sqlite_utils.Database(path)
-    db[table].add_column(col_name, col_type, fk=fk, fk_col=fk_col)
+    db[table].add_column(
+        col_name, col_type, fk=fk, fk_col=fk_col, not_null_default=not_null_default
+    )
 
 
 @cli.command(name="add-foreign-key")
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 07701bd..a77d9da 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -292,7 +292,9 @@ class Table:
         self.db.conn.execute(sql)
         return self
 
-    def add_column(self, col_name, col_type=None, fk=None, fk_col=None):
+    def add_column(
+        self, col_name, col_type=None, fk=None, fk_col=None, not_null_default=None
+    ):
         fk_col_type = None
         if fk is not None:
             # fk must be a valid table
@@ -313,10 +315,19 @@ class Table:
                     fk_col_type = "INTEGER"
         if col_type is None:
             col_type = str
-        sql = "ALTER TABLE [{table}] ADD COLUMN [{col_name}] {col_type};".format(
+        not_null_sql = None
+        if not_null_default is not None:
+            not_null_sql = "NOT NULL DEFAULT {}".format(
+                # Use SQLite itself to correctly escape this string:
+                self.db.conn.execute(
+                    "SELECT quote(:default)", {"default": not_null_default}
+                ).fetchone()[0]
+            )
+        sql = "ALTER TABLE [{table}] ADD COLUMN [{col_name}] {col_type}{not_null_default};".format(
             table=self.name,
             col_name=col_name,
             col_type=fk_col_type or COLUMN_TYPE_MAPPING[col_type],
+            not_null_default=(" " + not_null_sql) if not_null_sql else "",
         )
         self.db.conn.execute(sql)
         if fk is not None:
diff --git a/tests/.DS_Store b/tests/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6
GIT binary patch
literal 6148
zcmeH~Jr2S!425mzP>H1@V-^m;4Wg<&0T*E43hX&L&p$$qDprKhvt+--jT7}7np#A3
zem<@ulZcFPQ@L2!n>{z**<q8>++&mCkOWA81W14cNZ<zv;LbK1Poaz?KmsK2CSc!(
z0ynLxE!0092;Krf2c+FF_Fe*7ECH>lEfg7;MkzE(HCqgga^y>{tEnwC%0;vJ&^%eQ
zLs35+`xjp>T0<F0fCPF1$Cyrb|F7^5{eNG?83~ZUUlGt@xh*qZDeu<Z%US-OSsOPv
j)R!Z4KLME7ReXlK;d!wEw5GODWMKRea10D2@KpjYNUI8I

literal 0
HcmV?d00001

diff --git a/tests/test_cli.py b/tests/test_cli.py
index 484c517..f164635 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -199,6 +199,27 @@ def test_add_column(db_path, col_name, col_type, expected_schema):
     assert expected_schema == collapse_whitespace(db["dogs"].schema)
 
 
+def test_add_column_not_null_default(db_path):
+    db = Database(db_path)
+    db.create_table("dogs", {"name": str})
+    assert "CREATE TABLE [dogs] ( [name] TEXT )" == collapse_whitespace(
+        db["dogs"].schema
+    )
+    args = [
+        "add-column",
+        db_path,
+        "dogs",
+        "nickname",
+        "--not-null-default",
+        "dogs'dawg",
+    ]
+    assert 0 == CliRunner().invoke(cli.cli, args).exit_code
+    assert (
+        "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT NOT NULL DEFAULT 'dogs''dawg')"
+        == collapse_whitespace(db["dogs"].schema)
+    )
+
+
 def test_add_foreign_key(db_path):
     db = Database(db_path)
     db["authors"].insert_all(
diff --git a/tests/test_create.py b/tests/test_create.py
index c149764..b03e925 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -131,33 +131,47 @@ def test_create_error_if_invalid_foreign_keys(fresh_db):
 
 
 @pytest.mark.parametrize(
-    "col_name,col_type,expected_schema",
+    "col_name,col_type,not_null_default,expected_schema",
     (
-        ("nickname", str, "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT)"),
-        ("dob", datetime.date, "CREATE TABLE [dogs] ( [name] TEXT , [dob] TEXT)"),
-        ("age", int, "CREATE TABLE [dogs] ( [name] TEXT , [age] INTEGER)"),
-        ("weight", float, "CREATE TABLE [dogs] ( [name] TEXT , [weight] FLOAT)"),
-        ("text", "TEXT", "CREATE TABLE [dogs] ( [name] TEXT , [text] TEXT)"),
+        ("nickname", str, None, "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT)"),
+        ("dob", datetime.date, None, "CREATE TABLE [dogs] ( [name] TEXT , [dob] TEXT)"),
+        ("age", int, None, "CREATE TABLE [dogs] ( [name] TEXT , [age] INTEGER)"),
+        ("weight", float, None, "CREATE TABLE [dogs] ( [name] TEXT , [weight] FLOAT)"),
+        ("text", "TEXT", None, "CREATE TABLE [dogs] ( [name] TEXT , [text] TEXT)"),
         (
             "integer",
             "INTEGER",
+            None,
             "CREATE TABLE [dogs] ( [name] TEXT , [integer] INTEGER)",
         ),
-        ("float", "FLOAT", "CREATE TABLE [dogs] ( [name] TEXT , [float] FLOAT)"),
-        ("blob", "blob", "CREATE TABLE [dogs] ( [name] TEXT , [blob] BLOB)"),
+        ("float", "FLOAT", None, "CREATE TABLE [dogs] ( [name] TEXT , [float] FLOAT)"),
+        ("blob", "blob", None, "CREATE TABLE [dogs] ( [name] TEXT , [blob] BLOB)"),
         (
             "default_str",
             None,
+            None,
             "CREATE TABLE [dogs] ( [name] TEXT , [default_str] TEXT)",
         ),
+        (
+            "nickname",
+            str,
+            "",
+            "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT NOT NULL DEFAULT '')",
+        ),
+        (
+            "nickname",
+            str,
+            "dawg's dawg",
+            "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT NOT NULL DEFAULT 'dawg''s dawg')",
+        ),
     ),
 )
-def test_add_column(fresh_db, col_name, col_type, expected_schema):
+def test_add_column(fresh_db, col_name, col_type, not_null_default, expected_schema):
     fresh_db.create_table("dogs", {"name": str})
     assert "CREATE TABLE [dogs] ( [name] TEXT )" == collapse_whitespace(
         fresh_db["dogs"].schema
     )
-    fresh_db["dogs"].add_column(col_name, col_type)
+    fresh_db["dogs"].add_column(col_name, col_type, not_null_default=not_null_default)
     assert expected_schema == collapse_whitespace(fresh_db["dogs"].schema)