.add_foreign_keys() uses .transform() instead of PRAGMA writable_schema

Closes #577

This should solve all sorts of problems seen by users of platforms that throw errors on writable_schema.

Also added `add_foreign_keys=` and `foreign_keys=` parameters to `table.transform()`.

Author: simonw

docs/python-api.rst

@@ -1281,9 +1281,7 @@ Adding foreign key constraints
 
 The SQLite ``ALTER TABLE`` statement doesn't have the ability to add foreign key references to an existing column.
 
-It's possible to add these references through very careful manipulation of SQLite's ``sqlite_master`` table, using ``PRAGMA writable_schema``.
-
-``sqlite-utils`` can do this for you, though there is a significant risk of data corruption if something goes wrong so it is advisable to create a fresh copy of your database file before attempting this.
+The ``add_foreign_key()`` method here is a convenient wrapper around :ref:`table.transform() <python_api_transform>`.
 
 Here's an example of this mechanism in action:
 
@@ -1317,11 +1315,7 @@ To ignore the case where the key already exists, use ``ignore=True``:
 Adding multiple foreign key constraints at once
 -----------------------------------------------
 
-The final step in adding a new foreign key to a SQLite database is to run ``VACUUM``, to ensure the new foreign key is available in future introspection queries.
-
-``VACUUM`` against a large (multi-GB) database can take several minutes or longer. If you are adding multiple foreign keys using ``table.add_foreign_key(...)`` these can quickly add up.
-
-Instead, you can use ``db.add_foreign_keys(...)`` to add multiple foreign keys within a single transaction. This method takes a list of four-tuples, each one specifying a ``table``, ``column``, ``other_table`` and ``other_column``.
+You can use ``db.add_foreign_keys(...)`` to add multiple foreign keys in one go. This method takes a list of four-tuples, each one specifying a ``table``, ``column``, ``other_table`` and ``other_column``.
 
 Here's an example adding two foreign keys at once:
 
@@ -1388,6 +1382,8 @@ To keep the original table around instead of dropping it, pass the ``keep_table=
 
     table.transform(types={"age": int}, keep_table="original_table")
 
+.. _python_api_transform_alter_column_types:
+
 Altering column types
 ---------------------
 
@@ -1400,6 +1396,8 @@ To alter the type of a column, use the ``types=`` argument:
 
 See :ref:`python_api_add_column` for a list of available types.
 
+.. _python_api_transform_rename_columns:
+
 Renaming columns
 ----------------
 
@@ -1410,6 +1408,8 @@ The ``rename=`` parameter can rename columns:
     # Rename 'age' to 'initial_age':
     table.transform(rename={"age": "initial_age"})
 
+.. _python_api_transform_drop_columns:
+
 Dropping columns
 ----------------
 
@@ -1420,6 +1420,8 @@ To drop columns, pass them in the ``drop=`` set:
     # Drop the 'age' column:
     table.transform(drop={"age"})
 
+.. _python_api_transform_change_primary_keys:
+
 Changing primary keys
 ---------------------
 
@@ -1430,6 +1432,8 @@ To change the primary key for a table, use ``pk=``. This can be passed a single
     # Make `user_id` the new primary key
     table.transform(pk="user_id")
 
+.. _python_api_transform_change_not_null:
+
 Changing not null status
 ------------------------
 
@@ -1450,6 +1454,8 @@ If you want to take existing ``NOT NULL`` columns and change them to allow null
     # Make age allow NULL and switch weight to being NOT NULL:
     table.transform(not_null={"age": False, "weight": True})
 
+.. _python_api_transform_alter_column_defaults:
+
 Altering column defaults
 ------------------------
 
@@ -1463,6 +1469,8 @@ The ``defaults=`` parameter can be used to set or change the defaults for differ
     # Now remove the default from that column:
     table.transform(defaults={"age": None})
 
+.. _python_api_transform_change_column_order:
+
 Changing column order
 ---------------------
 
@@ -1473,6 +1481,45 @@ The ``column_order=`` parameter can be used to change the order of the columns.
     # Change column order
     table.transform(column_order=("name", "age", "id")
 
+.. _python_api_transform_add_foreign_key_constraints:
+
+Adding foreign key constraints
+------------------------------
+
+You can add one or more foreign key constraints to a table using the ``add_foreign_keys=`` parameter:
+
+.. code-block:: python
+
+    db["places"].transform(
+        add_foreign_keys=(
+            ("country", "country", "id"),
+            ("continent", "continent", "id")
+        )
+    )
+
+This accepts the same arguments described in :ref:`specifying foreign keys <python_api_foreign_keys>` - so you can specify them as a full tuple of ``(column, other_table, other_column)``, or you can take a shortcut and pass just the name of the column, provided the table can be automatically derived from the column name:
+
+.. code-block:: python
+
+    db["places"].transform(
+        add_foreign_keys=(("country", "continent"))
+    )
+
+.. _python_api_transform_replace_foreign_key_constraints:
+
+Replacing foreign key constraints
+---------------------------------
+
+The ``foreign_keys=`` parameter is similar to  to ``add_foreign_keys=`` but can be be used to replace all foreign key constraints on a table, dropping any that are not explicitly mentioned:
+
+.. code-block:: python
+
+    db["places"].transform(
+        add_foreign_keys=(("continent",))
+    )
+
+.. _python_api_transform_drop_foreign_key_constraints:
+
 Dropping foreign key constraints
 --------------------------------
 

sqlite_utils/db.py

@@ -156,14 +156,16 @@
 Trigger = namedtuple("Trigger", ("name", "table", "sql"))
 
 
-ForeignKeysType = Union[
-    Iterable[str],
-    Iterable[ForeignKey],
-    Iterable[Tuple[str, str]],
-    Iterable[Tuple[str, str, str]],
-    Iterable[Tuple[str, str, str, str]],
+ForeignKeyIndicator = Union[
+    str,
+    ForeignKey,
+    Tuple[str, str],
+    Tuple[str, str, str],
+    Tuple[str, str, str, str],
 ]
 
+ForeignKeysType = Union[Iterable[ForeignKeyIndicator], List[ForeignKeyIndicator]]
+
 
 class Default:
     pass
@@ -747,9 +749,16 @@ def execute_returning_dicts(
     def resolve_foreign_keys(
         self, name: str, foreign_keys: ForeignKeysType
     ) -> List[ForeignKey]:
-        # foreign_keys may be a list of column names, a list of ForeignKey tuples,
-        # a list of tuple-pairs or a list of tuple-triples. We want to turn
-        # it into a list of ForeignKey tuples
+        """
+        Given a list of differing foreign_keys definitions, return a list of
+        fully resolved ForeignKey() named tuples.
+
+        :param name: Name of table that foreign keys are being defined for
+        :param foreign_keys: List of foreign keys, each of which can be a
+            string, a ForeignKey() named tuple, a tuple of (column, other_table),
+            or a tuple of (column, other_table, other_column), or a tuple of
+            (table, column, other_table, other_column)
+        """
         table = cast(Table, self[name])
         if all(isinstance(fk, ForeignKey) for fk in foreign_keys):
             return cast(List[ForeignKey], foreign_keys)
@@ -767,12 +776,20 @@ def resolve_foreign_keys(
         ), "foreign_keys= should be a list of tuples"
         fks = []
         for tuple_or_list in foreign_keys:
+            if len(tuple_or_list) == 4:
+                assert (
+                    tuple_or_list[0] == name
+                ), "First item in {} should have been {}".format(tuple_or_list, name)
             assert len(tuple_or_list) in (
                 2,
                 3,
+                4,
             ), "foreign_keys= should be a list of tuple pairs or triples"
-            if len(tuple_or_list) == 3:
-                tuple_or_list = cast(Tuple[str, str, str], tuple_or_list)
+            if len(tuple_or_list) in (3, 4):
+                if len(tuple_or_list) == 4:
+                    tuple_or_list = cast(Tuple[str, str, str], tuple_or_list[1:])
+                else:
+                    tuple_or_list = cast(Tuple[str, str, str], tuple_or_list)
                 fks.append(
                     ForeignKey(
                         name, tuple_or_list[0], tuple_or_list[1], tuple_or_list[2]
@@ -864,7 +881,7 @@ def sort_key(p):
         for fk in foreign_keys:
             if fk.other_table == name and columns.get(fk.other_column):
                 continue
-            if not any(
+            if fk.other_column != "rowid" and not any(
                 c for c in self[fk.other_table].columns if c.name == fk.other_column
             ):
                 raise AlterError(
@@ -1148,32 +1165,14 @@ def add_foreign_keys(self, foreign_keys: Iterable[Tuple[str, str, str, str]]):
                     (table, column, other_table, other_column)
                 )
 
-        # Construct SQL for use with "UPDATE sqlite_master SET sql = ? WHERE name = ?"
-        table_sql: Dict[str, str] = {}
-        for table, column, other_table, other_column in foreign_keys_to_create:
-            old_sql = table_sql.get(table, self[table].schema)
-            extra_sql = ",\n   FOREIGN KEY([{column}]) REFERENCES [{other_table}]([{other_column}])\n".format(
-                column=column, other_table=other_table, other_column=other_column
-            )
-            # Stick that bit in at the very end just before the closing ')'
-            last_paren = old_sql.rindex(")")
-            new_sql = old_sql[:last_paren].strip() + extra_sql + old_sql[last_paren:]
-            table_sql[table] = new_sql
+        # Group them by table
+        by_table: Dict[str, List] = {}
+        for fk in foreign_keys_to_create:
+            by_table.setdefault(fk[0], []).append(fk)
+
+        for table, fks in by_table.items():
+            cast(Table, self[table]).transform(add_foreign_keys=fks)
 
-        # And execute it all within a single transaction
-        with self.conn:
-            cursor = self.conn.cursor()
-            schema_version = cursor.execute("PRAGMA schema_version").fetchone()[0]
-            cursor.execute("PRAGMA writable_schema = 1")
-            for table_name, new_sql in table_sql.items():
-                cursor.execute(
-                    "UPDATE sqlite_master SET sql = ? WHERE name = ?",
-                    (new_sql, table_name),
-                )
-            cursor.execute("PRAGMA schema_version = %d" % (schema_version + 1))
-            cursor.execute("PRAGMA writable_schema = 0")
-        # Have to VACUUM outside the transaction to ensure .foreign_keys property
-        # can see the newly created foreign key.
         self.vacuum()
 
     def index_foreign_keys(self):
@@ -1704,7 +1703,9 @@ def transform(
         pk: Optional[Any] = DEFAULT,
         not_null: Optional[Iterable[str]] = None,
         defaults: Optional[Dict[str, Any]] = None,
-        drop_foreign_keys: Optional[Iterable] = None,
+        drop_foreign_keys: Optional[Iterable[str]] = None,
+        add_foreign_keys: Optional[ForeignKeysType] = None,
+        foreign_keys: Optional[ForeignKeysType] = None,
         column_order: Optional[List[str]] = None,
         keep_table: Optional[str] = None,
     ) -> "Table":
@@ -1721,6 +1722,8 @@ def transform(
         :param not_null: Columns to set as ``NOT NULL``
         :param defaults: Default values for columns
         :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
+        :param add_foreign_keys: List of foreign keys to add to the table
+        :param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
         :param column_order: List of strings specifying a full or partial column order
           to use when creating the table
         :param keep_table: If specified, the existing table will be renamed to this and will not be
@@ -1735,6 +1738,8 @@ def transform(
             not_null=not_null,
             defaults=defaults,
             drop_foreign_keys=drop_foreign_keys,
+            add_foreign_keys=add_foreign_keys,
+            foreign_keys=foreign_keys,
             column_order=column_order,
             keep_table=keep_table,
         )
@@ -1765,6 +1770,8 @@ def transform_sql(
         not_null: Optional[Iterable[str]] = None,
         defaults: Optional[Dict[str, Any]] = None,
         drop_foreign_keys: Optional[Iterable] = None,
+        add_foreign_keys: Optional[ForeignKeysType] = None,
+        foreign_keys: Optional[ForeignKeysType] = None,
         column_order: Optional[List[str]] = None,
         tmp_suffix: Optional[str] = None,
         keep_table: Optional[str] = None,
@@ -1779,6 +1786,8 @@ def transform_sql(
         :param not_null: Columns to set as ``NOT NULL``
         :param defaults: Default values for columns
         :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
+        :param add_foreign_keys: List of foreign keys to add to the table
+        :param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
         :param column_order: List of strings specifying a full or partial column order
           to use when creating the table
         :param tmp_suffix: Suffix to use for the temporary table name
@@ -1788,6 +1797,45 @@ def transform_sql(
         types = types or {}
         rename = rename or {}
         drop = drop or set()
+
+        create_table_foreign_keys: List[ForeignKeyIndicator] = []
+
+        if foreign_keys is not None:
+            if add_foreign_keys is not None:
+                raise ValueError(
+                    "Cannot specify both foreign_keys and add_foreign_keys"
+                )
+            if drop_foreign_keys is not None:
+                raise ValueError(
+                    "Cannot specify both foreign_keys and drop_foreign_keys"
+                )
+            create_table_foreign_keys.extend(foreign_keys)
+        else:
+            # Construct foreign_keys from current, plus add_foreign_keys, minus drop_foreign_keys
+            create_table_foreign_keys = []
+            for table, column, other_table, other_column in self.foreign_keys:
+                # Copy over old foreign keys, unless we are dropping them
+                if (drop_foreign_keys is None) or (column not in drop_foreign_keys):
+                    create_table_foreign_keys.append(
+                        ForeignKey(
+                            table,
+                            rename.get(column) or column,
+                            other_table,
+                            other_column,
+                        )
+                    )
+            # Add new foreign keys
+            if add_foreign_keys is not None:
+                for fk in self.db.resolve_foreign_keys(self.name, add_foreign_keys):
+                    create_table_foreign_keys.append(
+                        ForeignKey(
+                            self.name,
+                            rename.get(fk.column) or fk.column,
+                            fk.other_table,
+                            fk.other_column,
+                        )
+                    )
+
         new_table_name = "{}_new_{}".format(
             self.name, tmp_suffix or os.urandom(6).hex()
         )
@@ -1847,14 +1895,6 @@ def transform_sql(
                 {rename.get(c) or c: v for c, v in defaults.items()}
             )
 
-        # foreign_keys
-        create_table_foreign_keys = []
-        for table, column, other_table, other_column in self.foreign_keys:
-            if (drop_foreign_keys is None) or (column not in drop_foreign_keys):
-                create_table_foreign_keys.append(
-                    (rename.get(column) or column, other_table, other_column)
-                )
-
         if column_order is not None:
             column_order = [rename.get(col) or col for col in column_order]