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 0000000..5008ddf Binary files /dev/null and b/tests/.DS_Store differ 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)