diff --git a/piccolo/apps/migrations/auto/schema_snapshot.py b/piccolo/apps/migrations/auto/schema_snapshot.py index 858c7ff6a..580280c52 100644 --- a/piccolo/apps/migrations/auto/schema_snapshot.py +++ b/piccolo/apps/migrations/auto/schema_snapshot.py @@ -4,6 +4,8 @@ from piccolo.apps.migrations.auto.diffable_table import DiffableTable from piccolo.apps.migrations.auto.migration_manager import MigrationManager +from piccolo.columns.column_types import ForeignKey +from piccolo.table import Table, create_table_class @dataclass @@ -49,6 +51,35 @@ def get_snapshot(self) -> list[DiffableTable]: table.tablename = rename_table.new_tablename break + # Foreign keys on other tables still point at the old table + # class, so the snapshot would keep diffing against the + # renamed model and emit an ``alter_column`` for each of them + # on every ``--auto`` run. + for table in tables: + for column in table.columns: + if not isinstance(column, ForeignKey): + continue + references = column._meta.params.get("references") + if not ( + isinstance(references, type) + and issubclass(references, Table) + and references.__name__ + == rename_table.old_class_name + and references._meta.tablename + == rename_table.old_tablename + ): + continue + renamed = create_table_class( + class_name=rename_table.new_class_name, + bases=(references,), + class_kwargs={ + "tablename": rename_table.new_tablename, + "schema": references._meta.schema, + }, + ) + column._meta.params["references"] = renamed + column._foreign_key_meta.references = renamed + for change_table_schema in manager.change_table_schemas: for table in tables: if table.tablename == change_table_schema.tablename: diff --git a/tests/apps/migrations/auto/test_schema_snapshot.py b/tests/apps/migrations/auto/test_schema_snapshot.py index fe9676cd7..f4af4eca8 100644 --- a/tests/apps/migrations/auto/test_schema_snapshot.py +++ b/tests/apps/migrations/auto/test_schema_snapshot.py @@ -1,7 +1,13 @@ from unittest import TestCase -from piccolo.apps.migrations.auto import MigrationManager, SchemaSnapshot +from piccolo.apps.migrations.auto import ( + DiffableTable, + MigrationManager, + SchemaSnapshot, +) +from piccolo.columns.column_types import ForeignKey from piccolo.constraints import Unique +from piccolo.table import Table class TestSchemaSnaphot(TestCase): @@ -64,6 +70,62 @@ def test_rename_table(self): self.assertTrue(snapshot[0].class_name == "Performer") self.assertTrue(snapshot[0].tablename == "performer") + def test_rename_table_updates_foreign_key_references(self): + """ + Renaming a table must also update the foreign keys on other tables + which point at it, otherwise the snapshot keeps diffing against the + renamed model and emits an ``alter_column`` on every ``--auto`` run. + """ + + class Band(Table, tablename="band"): + pass + + manager_1 = MigrationManager() + manager_1.add_table(class_name="Band", tablename="band") + manager_1.add_table(class_name="Concert", tablename="concert") + manager_1.add_column( + table_class_name="Concert", + tablename="concert", + column_name="band", + column_class_name="ForeignKey", + column_class=ForeignKey, + params={"references": Band}, + ) + + manager_2 = MigrationManager() + manager_2.rename_table( + old_class_name="Band", + old_tablename="band", + new_class_name="Performer", + new_tablename="performer", + ) + + schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2]) + concert = schema_snapshot.get_table_from_snapshot("Concert") + references = concert.columns[0]._meta.params["references"] + + self.assertEqual(references.__name__, "Performer") + self.assertEqual(references._meta.tablename, "performer") + + # The snapshot now agrees with the renamed model, so there is + # nothing left to migrate. + class Performer(Table, tablename="performer"): + pass + + class Concert(Table, tablename="concert"): + band = ForeignKey(Performer) + + delta = ( + DiffableTable( + class_name="Concert", + tablename="concert", + columns=[Concert.band], + ) + - concert + ) + + self.assertEqual(delta.alter_columns, []) + def test_add_column(self): """ Test adding columns.