Ability to add a column that is a foreign key reference

Python API:

    db["dogs"].add_column("species_id", fk="species")
    # or
    db["dogs"].add_column("species_id", fk="species", fk_col="ref")

CLI:

    $ sqlite-utils add-column mydb.db dogs species_id --fk species
    # or
    $ sqlite-utils add-column mydb.db dogs species_id --fk species --fk-col ref

Closes #16

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

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

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")

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(

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(

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",
     [