sqlite-utils transform --add-foreign-key option, closes #585

simonw · Aug 18, 2023 · 56093de · 56093de
1 parent 70717dc
commit 56093de
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 30 deletions.
diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
@@ -455,20 +455,25 @@ See :ref:`cli_transform_table`.
               --rename column2 column_renamed
 
     Options:
-      --type <TEXT CHOICE>...   Change column type to INTEGER, TEXT, FLOAT or BLOB
-      --drop TEXT               Drop this column
-      --rename <TEXT TEXT>...   Rename this column to X
-      -o, --column-order TEXT   Reorder columns
-      --not-null TEXT           Set this column to NOT NULL
-      --not-null-false TEXT     Remove NOT NULL from this column
-      --pk TEXT                 Make this column the primary key
-      --pk-none                 Remove primary key (convert to rowid table)
-      --default <TEXT TEXT>...  Set default value for this column
-      --default-none TEXT       Remove default from this column
-      --drop-foreign-key TEXT   Drop foreign key constraint for this column
-      --sql                     Output SQL without executing it
-      --load-extension TEXT     Path to SQLite extension, with optional :entrypoint
-      -h, --help                Show this message and exit.
+      --type <TEXT CHOICE>...         Change column type to INTEGER, TEXT, FLOAT or
+                                      BLOB
+      --drop TEXT                     Drop this column
+      --rename <TEXT TEXT>...         Rename this column to X
+      -o, --column-order TEXT         Reorder columns
+      --not-null TEXT                 Set this column to NOT NULL
+      --not-null-false TEXT           Remove NOT NULL from this column
+      --pk TEXT                       Make this column the primary key
+      --pk-none                       Remove primary key (convert to rowid table)
+      --default <TEXT TEXT>...        Set default value for this column
+      --default-none TEXT             Remove default from this column
+      --add-foreign-key <TEXT TEXT TEXT>...
+                                      Add a foreign key constraint from a column to
+                                      another table with another column
+      --drop-foreign-key TEXT         Drop foreign key constraint for this column
+      --sql                           Output SQL without executing it
+      --load-extension TEXT           Path to SQLite extension, with optional
+                                      :entrypoint
+      -h, --help                      Show this message and exit.
 
 
 .. _cli_ref_extract:

diff --git a/docs/cli.rst b/docs/cli.rst
@@ -2061,6 +2061,9 @@ Every option for this table (with the exception of ``--pk-none``) can be specifi
 ``--drop-foreign-key column``
     Drop the specified foreign key.
 
+``--add-foregn-key column other_table other_column``
+    Add a foreign key constraint to ``column`` pointing to ``other_table.other_column``.
+
 If you want to see the SQL that will be executed to make the change without actually executing it, add the ``--sql`` flag. For example:
 
 .. code-block:: bash

diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
@@ -2406,6 +2406,14 @@ def schema(
     "--default-none", type=str, multiple=True, help="Remove default from this column"
 )
 @click.option(
+    "add_foreign_keys",
+    "--add-foreign-key",
+    type=(str, str, str),
+    multiple=True,
+    help="Add a foreign key constraint from a column to another table with another column",
+)
+@click.option(
+    "drop_foreign_keys",
     "--drop-foreign-key",
     type=str,
     multiple=True,
@@ -2426,7 +2434,8 @@ def transform(
     pk_none,
     default,
     default_none,
-    drop_foreign_key,
+    add_foreign_keys,
+    drop_foreign_keys,
     sql,
     load_extension,
 ):
@@ -2475,8 +2484,10 @@ def transform(
     elif pk_none:
         kwargs["pk"] = None
     kwargs["defaults"] = default_dict
-    if drop_foreign_key:
-        kwargs["drop_foreign_keys"] = drop_foreign_key
+    if drop_foreign_keys:
+        kwargs["drop_foreign_keys"] = drop_foreign_keys
+    if add_foreign_keys:
+        kwargs["add_foreign_keys"] = add_foreign_keys
 
     if sql:
         for line in db[table].transform_sql(**kwargs):

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -1645,10 +1645,52 @@ def test_transform(db_path, args, expected_schema):
     assert schema == expected_schema
 
 
-def test_transform_drop_foreign_key(db_path):
+@pytest.mark.parametrize(
+    "extra_args,expected_schema",
+    (
+        (
+            ["--drop-foreign-key", "country"],
+            (
+                'CREATE TABLE "places" (\n'
+                "   [id] INTEGER PRIMARY KEY,\n"
+                "   [name] TEXT,\n"
+                "   [country] INTEGER,\n"
+                "   [city] INTEGER REFERENCES [city]([id]),\n"
+                "   [continent] INTEGER\n"
+                ")"
+            ),
+        ),
+        (
+            ["--drop-foreign-key", "country", "--drop-foreign-key", "city"],
+            (
+                'CREATE TABLE "places" (\n'
+                "   [id] INTEGER PRIMARY KEY,\n"
+                "   [name] TEXT,\n"
+                "   [country] INTEGER,\n"
+                "   [city] INTEGER,\n"
+                "   [continent] INTEGER\n"
+                ")"
+            ),
+        ),
+        (
+            ["--add-foreign-key", "continent", "continent", "id"],
+            (
+                'CREATE TABLE "places" (\n'
+                "   [id] INTEGER PRIMARY KEY,\n"
+                "   [name] TEXT,\n"
+                "   [country] INTEGER REFERENCES [country]([id]),\n"
+                "   [city] INTEGER REFERENCES [city]([id]),\n"
+                "   [continent] INTEGER REFERENCES [continent]([id])\n"
+                ")"
+            ),
+        ),
+    ),
+)
+def test_transform_add_or_drop_foreign_key(db_path, extra_args, expected_schema):
     db = Database(db_path)
     with db.conn:
         # Create table with three foreign keys so we can drop two of them
+        db["continent"].insert({"id": 1, "name": "Europe"}, pk="id")
         db["country"].insert({"id": 1, "name": "France"}, pk="id")
         db["city"].insert({"id": 24, "name": "Paris"}, pk="id")
         db["places"].insert(
@@ -1657,6 +1699,7 @@ def test_transform_drop_foreign_key(db_path):
                 "name": "Caveau de la Huchette",
                 "country": 1,
                 "city": 24,
+                "continent": 1,
             },
             foreign_keys=("country", "city"),
             pk="id",
@@ -1667,21 +1710,12 @@ def test_transform_drop_foreign_key(db_path):
             "transform",
             db_path,
             "places",
-            "--drop-foreign-key",
-            "country",
-        ],
+        ]
+        + extra_args,
     )
-    print(result.output)
     assert result.exit_code == 0
     schema = db["places"].schema
-    assert schema == (
-        'CREATE TABLE "places" (\n'
-        "   [id] INTEGER PRIMARY KEY,\n"
-        "   [name] TEXT,\n"
-        "   [country] INTEGER,\n"
-        "   [city] INTEGER REFERENCES [city]([id])\n"
-        ")"
-    )
+    assert schema == expected_schema
 
 
 _common_other_schema = (