diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index c988174..5e66454 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -722,12 +722,9 @@ class Table(Queryable): rename=None, drop=None, pk=None, - foreign_keys=None, - column_order=None, not_null=None, defaults=None, - hash_id=None, - extracts=None, + drop_foreign_keys=None, ): assert self.exists(), "Cannot transform a table that doesn't exist yet" sqls = self.transform_sql( @@ -737,10 +734,7 @@ class Table(Queryable): pk=pk, not_null=not_null, defaults=defaults, - # foreign_keys=foreign_keys, - # column_order=column_order, - # hash_id=hash_id, - # extracts=extracts, + drop_foreign_keys=drop_foreign_keys, ) with self.db.conn: for sql in sqls: @@ -755,6 +749,7 @@ class Table(Queryable): pk=None, not_null=None, defaults=None, + drop_foreign_keys=None, tmp_suffix=None, ): columns = columns or {} @@ -807,6 +802,16 @@ class Table(Queryable): {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, other_table, other_column) not in drop_foreign_keys + ): + create_table_foreign_keys.append( + (rename.get(column) or column, other_table, other_column) + ) + sqls.append( self.db.create_table_sql( new_table_name, @@ -814,9 +819,8 @@ class Table(Queryable): pk=pk, not_null=create_table_not_null, defaults=create_table_defaults, - # foreign_keys=foreign_keys, + foreign_keys=create_table_foreign_keys, # column_order=column_order, - # not_null=not_null, # hash_id=hash_id, # extracts=extracts, ).strip() diff --git a/tests/test_transform.py b/tests/test_transform.py index 6016033..7abbc8c 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -1,3 +1,4 @@ +from sqlite_utils.db import ForeignKey import pytest @@ -159,3 +160,84 @@ def test_remove_defaults(fresh_db): dogs.schema == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] INTEGER\n)' ) + + +@pytest.fixture +def authors_db(fresh_db): + books = fresh_db["books"] + authors = fresh_db["authors"] + authors.insert({"id": 5, "name": "Jane McGonical"}, pk="id") + books.insert( + {"id": 2, "title": "Reality is Broken", "author_id": 5}, + foreign_keys=("author_id",), + pk="id", + ) + return fresh_db + + +def test_transform_foreign_keys_persist(authors_db): + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", column="author_id", other_table="authors", other_column="id" + ) + ] + authors_db["books"].transform(rename={"title": "book_title"}) + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", column="author_id", other_table="authors", other_column="id" + ) + ] + + +def test_transform_foreign_keys_survive_renamed_column(authors_db): + authors_db["books"].transform(rename={"author_id": "author_id_2"}) + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", + column="author_id_2", + other_table="authors", + other_column="id", + ) + ] + + +def test_transform_drop_foreign_keys(fresh_db): + # Create table with three foreign keys so we can drop two of them + fresh_db["country"].insert({"id": 1, "name": "France"}, pk="id") + fresh_db["continent"].insert({"id": 2, "name": "Europe"}, pk="id") + fresh_db["city"].insert({"id": 24, "name": "Paris"}, pk="id") + fresh_db["places"].insert( + { + "id": 32, + "name": "Caveau de la Huchette", + "country": 1, + "continent": 2, + "city": 24, + }, + foreign_keys=("country", "continent", "city"), + ) + assert fresh_db["places"].foreign_keys == [ + ForeignKey( + table="places", column="city", other_table="city", other_column="id" + ), + ForeignKey( + table="places", + column="continent", + other_table="continent", + other_column="id", + ), + ForeignKey( + table="places", column="country", other_table="country", other_column="id" + ), + ] + # Drop two of those foreign keys + fresh_db["places"].transform( + drop_foreign_keys=( + ("country", "country", "id"), + ("continent", "continent", "id"), + ) + ) + # Should be only one foreign key now + assert fresh_db["places"].foreign_keys == [ + ForeignKey(table="places", column="city", other_table="city", other_column="id") + ]