fix table creation ordering in migrations (piccolo-orm#211)

* fix table creation ordering

* put a recursion limit on `_compare_tables`

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -3,6 +3,7 @@
 import inspect
 import typing as t
 from dataclasses import dataclass, field
+from functools import cmp_to_key
 
 from piccolo.apps.migrations.auto.diffable_table import DiffableTable
 from piccolo.apps.migrations.auto.operations import (
@@ -117,6 +118,49 @@ def table_class_names(self) -> t.List[str]:
         return list(set([i.table_class_name for i in self.alter_columns]))
 
 
+def _compare_tables(
+    table_a: t.Type[Table],
+    table_b: t.Type[Table],
+    iterations: int = 0,
+    max_iterations=5,
+) -> int:
+    """
+    A comparison function, for sorting Table classes, based on their foreign
+    keys.
+
+    :param iterations:
+        As this function is called recursively, we use this to limit the depth,
+        to prevent an infinite loop.
+
+    """
+    if iterations >= max_iterations:
+        return 0
+
+    for fk_column in table_a._meta.foreign_key_columns:
+        references = fk_column._foreign_key_meta.resolved_references
+        if references._meta.tablename == table_b._meta.tablename:
+            return 1
+        else:
+            for _fk_column in references._meta.foreign_key_columns:
+                _references = _fk_column._foreign_key_meta.resolved_references
+                if _compare_tables(
+                    _references, table_b, iterations=iterations + 1
+                ):
+                    return 1
+
+    return -1
+
+
+def sort_table_classes(
+    table_classes: t.List[t.Type[Table]],
+) -> t.List[t.Type[Table]]:
+    """
+    Sort the table classes based on their foreign keys, so they can be created
+    in the correct order.
+    """
+    return sorted(table_classes, key=cmp_to_key(_compare_tables))
+
+
 @dataclass
 class MigrationManager:
     """
@@ -563,25 +607,29 @@ async def _run_rename_columns(self, backwards=False):
                 ).run()
 
     async def _run_add_tables(self, backwards=False):
+        table_classes: t.List[t.Type[Table]] = []
+        for add_table in self.add_tables:
+            add_columns: t.List[
+                AddColumnClass
+            ] = self.add_columns.for_table_class_name(add_table.class_name)
+            _Table: t.Type[Table] = create_table_class(
+                class_name=add_table.class_name,
+                class_kwargs={"tablename": add_table.tablename},
+                class_members={
+                    add_column.column._meta.name: add_column.column
+                    for add_column in add_columns
+                },
+            )
+            table_classes.append(_Table)
+
+        # Sort by foreign key, so they're created in the right order.
+        sorted_table_classes = sort_table_classes(table_classes)
+
         if backwards:
-            for add_table in self.add_tables:
-                await add_table.to_table_class().alter().drop_table(
-                    cascade=True
-                ).run()
+            for _Table in reversed(sorted_table_classes):
+                await _Table.alter().drop_table(cascade=True).run()
         else:
-            for add_table in self.add_tables:
-                add_columns: t.List[
-                    AddColumnClass
-                ] = self.add_columns.for_table_class_name(add_table.class_name)
-                _Table: t.Type[Table] = create_table_class(
-                    class_name=add_table.class_name,
-                    class_kwargs={"tablename": add_table.tablename},
-                    class_members={
-                        add_column.column._meta.name: add_column.column
-                        for add_column in add_columns
-                    },
-                )
-
+            for _Table in sorted_table_classes:
                 await _Table.create_table().run()
 
     async def _run_add_columns(self, backwards=False):

tests/apps/migrations/auto/test_migration_manager.py

@@ -1,19 +1,40 @@
 import asyncio
+from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
-from piccolo.apps.migrations.auto import MigrationManager
+from piccolo.apps.migrations.auto.migration_manager import (
+    MigrationManager,
+    sort_table_classes,
+)
 from piccolo.apps.migrations.commands.base import BaseMigrationManager
 from piccolo.columns import Text, Varchar
 from piccolo.columns.base import OnDelete, OnUpdate
 from piccolo.columns.column_types import ForeignKey
 from piccolo.conf.apps import AppConfig
 from piccolo.utils.lazy_loader import LazyLoader
 from tests.base import DBTestCase, postgres_only, set_mock_return_value
-from tests.example_app.tables import Manager
+from tests.example_app.tables import Band, Concert, Manager, Venue
 
 asyncpg = LazyLoader("asyncpg", globals(), "asyncpg")
 
 
+class TestSortTableClasses(TestCase):
+    def test_sort_table_classes(self):
+        self.assertEqual(sort_table_classes([Manager, Band]), [Manager, Band])
+        self.assertEqual(sort_table_classes([Band, Manager]), [Manager, Band])
+
+        sorted_tables = sort_table_classes([Manager, Venue, Concert, Band])
+        self.assertTrue(
+            sorted_tables.index(Manager) < sorted_tables.index(Band)
+        )
+        self.assertTrue(
+            sorted_tables.index(Venue) < sorted_tables.index(Concert)
+        )
+        self.assertTrue(
+            sorted_tables.index(Band) < sorted_tables.index(Concert)
+        )
+
+
 class TestMigrationManager(DBTestCase):
     @postgres_only
     def test_rename_column(self):

tests/example_app/piccolo_migrations/2020-12-17T18-44-30.py

@@ -14,8 +14,8 @@ class Manager(Table, tablename="manager"):
 async def forwards():
     manager = MigrationManager(migration_id=ID, app_name="example_app")
 
-    manager.add_table("Manager", tablename="manager")
     manager.add_table("Band", tablename="band")
+    manager.add_table("Manager", tablename="manager")
 
     manager.add_column(
         table_class_name="Band",