Schema generation fixes (piccolo-orm#226)

* fix a typo in the column import

* sort the tables by their foreign keys

* make sure there are no missing imports from column params

* check the generated code is valid

* use graphlib for table sorting

Using just basic sort would fail on occasion

* remove walrus operators

* remove unnecessary imports from schema generation output

* make graph creation recursive

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -3,7 +3,6 @@
 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 (
@@ -16,6 +15,7 @@
 from piccolo.columns import Column, column_types
 from piccolo.engine import engine_finder
 from piccolo.table import Table, create_table_class
+from piccolo.utils.graphlib import TopologicalSorter
 
 
 @dataclass
@@ -118,37 +118,48 @@ 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],
+def _get_graph(
+    table_classes: t.List[t.Type[Table]],
     iterations: int = 0,
-    max_iterations=5,
-) -> int:
+    max_iterations: int = 5,
+) -> t.Dict[str, t.Set[str]]:
     """
-    A comparison function, for sorting Table classes, based on their foreign
-    keys.
+    Analyses the tables based on their foreign keys, and returns a data
+    structure like:
+
+    .. code-block:: python
 
-    :param iterations:
-        As this function is called recursively, we use this to limit the depth,
-        to prevent an infinite loop.
+        {'band': {'manager'}, 'concert': {'band', 'venue'}, 'manager': set()}
+
+    The keys are tablenames, and the values are tablenames directly connected
+    to it via a foreign key.
 
     """
+    output: t.Dict[str, t.Set[str]] = {}
+
     if iterations >= max_iterations:
-        return 0
+        return output
 
-    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
+    for table_class in table_classes:
+        dependents: t.Set[str] = set()
+        for fk in table_class._meta.foreign_key_columns:
+            dependents.add(
+                fk._foreign_key_meta.resolved_references._meta.tablename
+            )
 
-    return -1
+            # We also recursively check the related tables to get a fuller
+            # picture of the schema and relationships.
+            referenced_table = fk._foreign_key_meta.resolved_references
+            output.update(
+                _get_graph(
+                    [referenced_table],
+                    iterations=iterations + 1,
+                )
+            )
+
+        output[table_class._meta.tablename] = dependents
+
+    return output
 
 
 def sort_table_classes(
@@ -158,7 +169,23 @@ def sort_table_classes(
     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))
+    table_class_dict = {
+        table_class._meta.tablename: table_class
+        for table_class in table_classes
+    }
+
+    graph = _get_graph(table_classes)
+
+    sorter = TopologicalSorter(graph)
+    ordered_tablenames = tuple(sorter.static_order())
+
+    output: t.List[t.Type[Table]] = []
+    for tablename in ordered_tablenames:
+        table_class = table_class_dict.get(tablename, None)
+        if table_class is not None:
+            output.append(table_class)
+
+    return output
 
 
 @dataclass

piccolo/apps/schema/commands/generate.py

@@ -6,6 +6,8 @@
 import black
 from typing_extensions import Literal
 
+from piccolo.apps.migrations.auto.migration_manager import sort_table_classes
+from piccolo.apps.migrations.auto.serialisation import serialise_params
 from piccolo.columns.base import Column
 from piccolo.columns.column_types import (
     JSON,
@@ -336,18 +338,22 @@ class Schema(Table, db=engine):
                         )
                     else:
                         kwargs["references"] = ForeignKeyPlaceholder
-                    imports.add(
-                        "from piccolo.columns.base import OnDelete, OnUpdate"
-                    )
 
                 imports.add(
-                    "from piccolo.column_types import " + column_type.__name__
+                    "from piccolo.columns.column_types import "
+                    + column_type.__name__
                 )
 
                 if column_type is Varchar:
                     kwargs["length"] = pg_row_meta.character_maximum_length
 
-                columns[column_name] = column_type(**kwargs)
+                column = column_type(**kwargs)
+
+                serialised_params = serialise_params(column._meta.params)
+                for extra_import in serialised_params.extra_imports:
+                    imports.add(extra_import.__repr__())
+
+                columns[column_name] = column
             else:
                 warnings.append(f"{tablename}.{column_name} ['{data_type}']")
 
@@ -358,6 +364,13 @@ class Schema(Table, db=engine):
         )
         tables.append(table)
 
+    # Sort the tables based on their ForeignKeys.
+    tables = sort_table_classes(tables)
+
+    # We currently don't show the index argument for columns in the output,
+    # so we don't need this import for now:
+    imports.remove("from piccolo.columns.indexes import IndexMethod")
+
     return OutputSchema(
         imports=sorted(list(imports)), warnings=warnings, tables=tables
     )

piccolo/utils/graphlib/__init__.py

@@ -0,0 +1,5 @@
+try:
+    from graphlib import CycleError, TopologicalSorter  # type: ignore
+except ImportError:
+    # For version < Python 3.9
+    from ._graphlib import CycleError, TopologicalSorter