diff --git a/docs/cli.rst b/docs/cli.rst index 0857d87..8945920 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -266,6 +266,16 @@ You can add a column using the ``add-column`` command:: The last argument here is the type of the column to be created. You can use one of ``text``, ``integer``, ``float`` or ``blob``. If you leave it off, ``text`` will be used. +You can add a column that is a foreign key reference to another table using the ``--fk`` option:: + + $ sqlite-utils add-column mydb.db dogs species_id --fk species + +This will automatically detect the name of the primary key on the species table and use that (and its type) for the new column. + +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 + .. _cli_add_column_alter: Adding columns automatically on insert/update diff --git a/docs/python-api.rst b/docs/python-api.rst index 89f510c..90cc159 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -266,6 +266,20 @@ If you pass a Python type, it will be mapped to SQLite types as shown here:: np.float32: "FLOAT" np.float64: "FLOAT" +You can also add a column that is a foreign key reference to another table using the ``fk`` parameter: + +.. code-block:: python + + db["dogs"].add_column("species_id", fk="species") + +This will automatically detect the name of the primary key on the species table and use that (and its type) for the new column. + +You can explicitly specify the column you wish to reference using ``fk_col``: + +.. code-block:: python + + db["dogs"].add_column("species_id", fk="species", fk_col="ref") + .. _python_api_add_column_alter: Adding columns automatically on insert/update diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 531d017..daf4abc 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -169,10 +169,12 @@ def optimize(path, no_vacuum): ), required=False, ) -def add_column(path, table, col_name, col_type): +@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): "Add a column to the specified table" db = sqlite_utils.Database(path) - db[table].add_column(col_name, col_type) + db[table].add_column(col_name, col_type, fk=fk, fk_col=fk_col) @cli.command(name="add-foreign-key") diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index b354d22..07701bd 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -292,13 +292,35 @@ class Table: self.db.conn.execute(sql) return self - def add_column(self, col_name, col_type=None): + def add_column(self, col_name, col_type=None, fk=None, fk_col=None): + fk_col_type = None + if fk is not None: + # fk must be a valid table + if not fk in self.db.table_names(): + raise AlterError("table '{}' does not exist".format(fk)) + # if fk_col specified, must be a valid column + if fk_col is not None: + if fk_col not in self.db[fk].columns_dict: + raise AlterError("table '{}' has no column {}".format(fk, fk_col)) + else: + # automatically set fk_col to first primary_key of fk table + pks = [c for c in self.db[fk].columns if c.is_pk] + if pks: + fk_col = pks[0].name + fk_col_type = pks[0].type + else: + fk_col = "rowid" + fk_col_type = "INTEGER" if col_type is None: col_type = str sql = "ALTER TABLE [{table}] ADD COLUMN [{col_name}] {col_type};".format( - table=self.name, col_name=col_name, col_type=COLUMN_TYPE_MAPPING[col_type] + table=self.name, + col_name=col_name, + col_type=fk_col_type or COLUMN_TYPE_MAPPING[col_type], ) self.db.conn.execute(sql) + if fk is not None: + self.add_foreign_key(col_name, fk, fk_col) return self def drop(self): @@ -306,7 +328,10 @@ class Table: def add_foreign_key(self, column, other_table, other_column): # Sanity check that the other column exists - if not [c for c in self.db[other_table].columns if c.name == other_column]: + if ( + not [c for c in self.db[other_table].columns if c.name == other_column] + and other_column != "rowid" + ): raise AlterError("No such column: {}.{}".format(other_table, other_column)) # Check we do not already have an existing foreign key if any( diff --git a/tests/test_cli.py b/tests/test_cli.py index ec30c0b..484c517 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -233,7 +233,7 @@ def test_add_foreign_key(db_path): "Error: Foreign key already exists for author_id => authors.id" == result.output.strip() ) - # Error if we try against an invalid cgolumn + # Error if we try against an invalid column result = CliRunner().invoke( cli.cli, ["add-foreign-key", db_path, "books", "author_id", "authors", "bad"] ) @@ -241,6 +241,48 @@ def test_add_foreign_key(db_path): assert "Error: No such column: authors.bad" == result.output.strip() +def test_add_column_foreign_key(db_path): + db = Database(db_path) + db["authors"].insert({"id": 1, "name": "Sally"}, pk="id") + db["books"].insert({"title": "Hedgehogs of the world"}) + # Add an author_id foreign key column to the books table + result = CliRunner().invoke( + cli.cli, ["add-column", db_path, "books", "author_id", "--fk", "authors"] + ) + assert 0 == result.exit_code, result.output + assert ( + "CREATE TABLE [books] ( [title] TEXT , [author_id] INTEGER, FOREIGN KEY(author_id) REFERENCES authors(id) )" + == collapse_whitespace(db["books"].schema) + ) + # Try it again with a custom --fk-col + result = CliRunner().invoke( + cli.cli, + [ + "add-column", + db_path, + "books", + "author_name_ref", + "--fk", + "authors", + "--fk-col", + "name", + ], + ) + assert 0 == result.exit_code, result.output + assert ( + "CREATE TABLE [books] ( [title] TEXT , [author_id] INTEGER, [author_name_ref] TEXT, " + "FOREIGN KEY(author_id) REFERENCES authors(id), " + "FOREIGN KEY(author_name_ref) REFERENCES authors(name) )" + == collapse_whitespace(db["books"].schema) + ) + # Throw an error if the --fk table does not exist + result = CliRunner().invoke( + cli.cli, ["add-column", db_path, "books", "author_id", "--fk", "bobcats"] + ) + assert 0 != result.exit_code + assert "table 'bobcats' does not exist" in str(result.exception) + + def test_enable_fts(db_path): assert None == Database(db_path)["Gosh"].detect_fts() result = CliRunner().invoke( diff --git a/tests/test_create.py b/tests/test_create.py index 227d4dd..c149764 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -195,6 +195,25 @@ def test_add_foreign_key_error_if_already_exists(fresh_db): assert "Foreign key already exists for author_id => authors.id" == ex.value.args[0] +def test_add_column_foreign_key(fresh_db): + fresh_db.create_table("dogs", {"name": str}) + fresh_db.create_table("breeds", {"name": str}) + fresh_db["dogs"].add_column("breed_id", fk="breeds") + assert ( + "CREATE TABLE [dogs] ( [name] TEXT , [breed_id] INTEGER, FOREIGN KEY(breed_id) REFERENCES breeds(rowid) )" + == collapse_whitespace(fresh_db["dogs"].schema) + ) + # And again with an explicit primary key column + fresh_db.create_table("subbreeds", {"name": str, "primkey": str}, pk="primkey") + fresh_db["dogs"].add_column("subbreed_id", fk="subbreeds") + assert ( + "CREATE TABLE [dogs] ( [name] TEXT , [breed_id] INTEGER, [subbreed_id] TEXT, " + "FOREIGN KEY(breed_id) REFERENCES breeds(rowid), " + "FOREIGN KEY(subbreed_id) REFERENCES subbreeds(primkey) )" + == collapse_whitespace(fresh_db["dogs"].schema) + ) + + @pytest.mark.parametrize( "extra_data,expected_new_columns", [