Merge pull request #53 from piercefreeman/bug/fix-foreign-key-dependencies

Add destination table/col to migration dependencies

Author: piercefreeman

iceaxe/__tests__/schemas/test_db_memory_serializer.py

@@ -1314,3 +1314,115 @@ class AutoAssignedModel(TableBase):
     )
     assert not auto_id_column.nullable
     assert auto_id_column.autoincrement
+
+
+@pytest.mark.asyncio
+async def test_foreign_key_table_dependency():
+    """
+    Test that foreign key constraints properly depend on the referenced table being created first.
+    This test verifies that the foreign key constraint is ordered after both tables are created.
+    """
+
+    class TargetModel(TableBase):
+        id: int = Field(primary_key=True)
+        value: str
+
+    class SourceModel(TableBase):
+        id: int = Field(primary_key=True)
+        target_id: int = Field(foreign_key="targetmodel.id")
+
+    migrator = DatabaseMemorySerializer()
+
+    # Make sure Source is parsed before Target so we can make sure our foreign-key
+    # constraint actually re-orders the final objects.
+    db_objects = list(migrator.delegate([SourceModel, TargetModel]))
+    ordering = migrator.order_db_objects(db_objects)
+
+    # Get all objects in their sorted order
+    sorted_objects = sorted(
+        [obj for obj, _ in db_objects], key=lambda obj: ordering[obj]
+    )
+
+    # Find the positions of key objects
+    target_table_pos = next(
+        i
+        for i, obj in enumerate(sorted_objects)
+        if isinstance(obj, DBTable) and obj.table_name == "targetmodel"
+    )
+    source_table_pos = next(
+        i
+        for i, obj in enumerate(sorted_objects)
+        if isinstance(obj, DBTable) and obj.table_name == "sourcemodel"
+    )
+    target_column_pos = next(
+        i
+        for i, obj in enumerate(sorted_objects)
+        if isinstance(obj, DBColumn)
+        and obj.table_name == "targetmodel"
+        and obj.column_name == "id"
+    )
+    target_pk_pos = next(
+        i
+        for i, obj in enumerate(sorted_objects)
+        if isinstance(obj, DBConstraint)
+        and obj.constraint_type == ConstraintType.PRIMARY_KEY
+        and obj.table_name == "targetmodel"
+    )
+    fk_constraint_pos = next(
+        i
+        for i, obj in enumerate(sorted_objects)
+        if isinstance(obj, DBConstraint)
+        and obj.constraint_type == ConstraintType.FOREIGN_KEY
+        and obj.table_name == "sourcemodel"
+    )
+
+    # The foreign key constraint should come after both tables and the target column are created
+    assert (
+        target_table_pos < fk_constraint_pos
+    ), "Foreign key constraint should be created after target table"
+    assert (
+        source_table_pos < fk_constraint_pos
+    ), "Foreign key constraint should be created after source table"
+    assert (
+        target_column_pos < fk_constraint_pos
+    ), "Foreign key constraint should be created after target column"
+    assert (
+        target_pk_pos < fk_constraint_pos
+    ), "Foreign key constraint should be created after target primary key"
+
+    # Verify the actual migration actions
+    actor = DatabaseActions()
+    actions = await migrator.build_actions(
+        actor, [], {}, [obj for obj, _ in db_objects], ordering
+    )
+
+    # Extract the table creation and foreign key constraint actions
+    table_creations = [
+        action
+        for action in actions
+        if isinstance(action, DryRunAction) and action.fn == actor.add_table
+    ]
+    fk_constraints = [
+        action
+        for action in actions
+        if isinstance(action, DryRunAction)
+        and action.fn == actor.add_constraint
+        and action.kwargs.get("constraint") == ConstraintType.FOREIGN_KEY
+    ]
+
+    # Verify that table creations come before foreign key constraints
+    assert len(table_creations) == 2
+    assert len(fk_constraints) == 1
+
+    table_creation_indices = [
+        i for i, action in enumerate(actions) if action in table_creations
+    ]
+    fk_constraint_indices = [
+        i for i, action in enumerate(actions) if action in fk_constraints
+    ]
+
+    assert all(
+        table_idx < fk_idx
+        for table_idx in table_creation_indices
+        for fk_idx in fk_constraint_indices
+    )

iceaxe/migrations/action_sorter.py

@@ -20,7 +20,7 @@ def __init__(self, graph: dict[DBObject, list[DBObject]]):
         self.in_degree = defaultdict(int)
         self.nodes = set(graph.keys())
 
-        for node, dependencies in graph.items():
+        for node, dependencies in list(graph.items()):
             for dep in dependencies:
                 self.in_degree[node] += 1
                 if dep not in self.nodes:

iceaxe/migrations/cli.py

@@ -4,9 +4,6 @@
 from iceaxe.base import DBModelMetaclass, TableBase
 from iceaxe.io import resolve_package_path
 from iceaxe.logging import CONSOLE
-from iceaxe.migrations.client_io import fetch_migrations, sort_migrations
-from iceaxe.migrations.generator import MigrationGenerator
-from iceaxe.migrations.migrator import Migrator
 from iceaxe.schemas.db_serializer import DatabaseSerializer
 from iceaxe.session import DBConnection
 
@@ -35,6 +32,11 @@ def generate_migration(message: str):
         handle_generate("my_project", db_connection, message=message)
     ```
     """
+    # Any local imports must be done here to avoid circular imports because migrations.__init__
+    # imports this file.
+    from iceaxe.migrations.client_io import fetch_migrations
+    from iceaxe.migrations.generator import MigrationGenerator
+    from iceaxe.migrations.migrator import Migrator
 
     CONSOLE.print("[bold blue]Generating migration to current schema")
 
@@ -122,6 +124,8 @@ async def handle_apply(
         project that's specified in pyproject.toml or setup.py.
 
     """
+    from iceaxe.migrations.client_io import fetch_migrations, sort_migrations
+    from iceaxe.migrations.migrator import Migrator
 
     migrations_path = resolve_package_path(package) / "migrations"
     if not migrations_path.exists():
@@ -178,6 +182,8 @@ async def handle_rollback(
         project that's specified in pyproject.toml or setup.py.
 
     """
+    from iceaxe.migrations.client_io import fetch_migrations, sort_migrations
+    from iceaxe.migrations.migrator import Migrator
 
     migrations_path = resolve_package_path(package) / "migrations"
     if not migrations_path.exists():

iceaxe/schemas/db_memory_serializer.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from datetime import date, datetime, time, timedelta
 from inspect import isgenerator
-from typing import Any, Generator, Sequence, Type, TypeVar
+from typing import Any, Generator, Sequence, Type, TypeVar, Union
 from uuid import UUID
 
 from pydantic_core import PydanticUndefined
@@ -29,9 +29,12 @@
 )
 from iceaxe.schemas.db_stubs import (
     DBColumn,
+    DBColumnPointer,
     DBConstraint,
+    DBConstraintPointer,
     DBObject,
     DBObjectPointer,
+    DBPointerOr,
     DBTable,
     DBType,
     DBTypePointer,
@@ -44,11 +47,13 @@
     PRIMITIVE_WRAPPER_TYPES,
 )
 
+NodeYieldType = Union[DBObject, DBObjectPointer, "NodeDefinition"]
+
 
 @dataclass
 class NodeDefinition:
     node: DBObject
-    dependencies: list[DBObject]
+    dependencies: list[DBObject | DBObjectPointer]
     force_no_dependencies: bool
 
 
@@ -110,18 +115,42 @@ def order_db_objects(
         for _, dependencies in db_objects:
             for dep in dependencies:
                 if isinstance(dep, DBObjectPointer):
-                    if dep.representation() not in db_objects_by_name:
+                    if isinstance(dep, DBPointerOr):
+                        # For OR pointers, at least one of the pointers must be resolvable
+                        if not any(
+                            pointer.representation() in db_objects_by_name
+                            for pointer in dep.pointers
+                        ):
+                            raise ValueError(
+                                f"None of the OR pointers {[p.representation() for p in dep.pointers]} found in the defined database objects"
+                            )
+                    elif dep.representation() not in db_objects_by_name:
                         raise ValueError(
                             f"Pointer {dep.representation()} not found in the defined database objects"
                         )
 
         # Map the potentially different objects to the same object
-        graph_edges = {
-            db_objects_by_name[obj.representation()]: [
-                db_objects_by_name[dep.representation()] for dep in dependencies
-            ]
-            for obj, dependencies in db_objects
-        }
+        graph_edges = {}
+        for obj, dependencies in db_objects:
+            resolved_deps = []
+            for dep in dependencies:
+                if isinstance(dep, DBObjectPointer):
+                    if isinstance(dep, DBPointerOr):
+                        # Add all resolvable pointers as dependencies
+                        resolved_deps.extend(
+                            db_objects_by_name[pointer.representation()]
+                            for pointer in dep.pointers
+                            if pointer.representation() in db_objects_by_name
+                        )
+                    else:
+                        resolved_deps.append(db_objects_by_name[dep.representation()])
+                else:
+                    resolved_deps.append(dep)
+
+            if isinstance(obj, DBObjectPointer):
+                continue
+
+            graph_edges[db_objects_by_name[obj.representation()]] = resolved_deps
 
         # Construct the directed acyclic graph
         ts = ActionTopologicalSorter(graph_edges)
@@ -217,9 +246,6 @@ def migrate(self, previous, actor: DatabaseActions):
         raise NotImplementedError()
 
 
-NodeYieldType = DBObject | NodeDefinition
-
-
 class DatabaseHandler:
     def __init__(self):
         self.python_to_sql = {
@@ -430,7 +456,32 @@ def _build_constraint(
                         target_table=target_table,
                         target_columns=frozenset({target_column}),
                     ),
-                )
+                ),
+                dependencies=[
+                    # Additional dependencies to ensure the target table/column is created first
+                    DBTable(table_name=target_table),
+                    DBColumnPointer(
+                        table_name=target_table,
+                        column_name=target_column,
+                    ),
+                    # Ensure the primary key constraint exists before the foreign key
+                    # constraint. Postgres also accepts a unique constraint on the same.
+                    DBPointerOr(
+                        pointers=tuple(
+                            [
+                                DBConstraintPointer(
+                                    table_name=target_table,
+                                    columns=frozenset([target_column]),
+                                    constraint_type=constraint_type,
+                                )
+                                for constraint_type in [
+                                    ConstraintType.PRIMARY_KEY,
+                                    ConstraintType.UNIQUE,
+                                ]
+                            ]
+                        ),
+                    ),
+                ],
             )
 
         if info.index:
@@ -509,10 +560,10 @@ def _yield_nodes(
         """
 
         def _format_dependencies(dependencies: Sequence[NodeYieldType]):
-            all_dependencies: list[DBObject] = []
+            all_dependencies: list[DBObject | DBObjectPointer] = []
 
             for value in dependencies:
-                if isinstance(value, DBObject):
+                if isinstance(value, (DBObject, DBObjectPointer)):
                     all_dependencies.append(value)
                 elif isinstance(value, NodeDefinition):
                     all_dependencies.append(value.node)

iceaxe/schemas/db_stubs.py

@@ -374,3 +374,32 @@ def merge(self, other: "DBType") -> "DBType":
             values=self.values,
             reference_columns=self.reference_columns | other.reference_columns,
         )
+
+
+class DBConstraintPointer(DBObjectPointer):
+    """
+    A pointer to a constraint that will be created. Used for dependency tracking
+    without needing to know the full constraint definition.
+    """
+
+    table_name: str
+    columns: frozenset[str]
+    constraint_type: ConstraintType
+
+    def representation(self) -> str:
+        # Match the representation of DBConstraint
+        return f"{self.table_name}.{sorted(self.columns)}.{self.constraint_type}"
+
+
+class DBPointerOr(DBObjectPointer):
+    """
+    A pointer that represents an OR relationship between multiple pointers.
+    When resolving dependencies, any of the provided pointers being present
+    will satisfy the dependency.
+    """
+
+    pointers: tuple[DBObjectPointer, ...]
+
+    def representation(self) -> str:
+        # Sort the representations to ensure consistent ordering
+        return "OR(" + ",".join(sorted(p.representation() for p in self.pointers)) + ")"