Foreign key migrations

py-modules/cli/macrostrat/cli/schema_management/__init__.py

@@ -5,24 +5,23 @@
 from typing import Callable
 
 import click
+from macrostrat.app_frame import CommandBase
+from macrostrat.database import Database
+from macrostrat.database.transfer import pg_dump_to_file
+from macrostrat.database.transfer.utils import raw_database_url
+from macrostrat.utils import get_logger
+from macrostrat.utils.shell import run
 from results import db as results_db
 from results.dbdiff import Migration
 from rich import print
+from sqlparse import format as format_sql
 from typer import Argument, Option
 
-from macrostrat.app_frame import CommandBase
 from macrostrat.core import app as macrostrat_app
 from macrostrat.core.config import settings
 from macrostrat.core.database import get_database
 from macrostrat.core.exc import MacrostratError
 from macrostrat.core.migrations import run_migrations
-from macrostrat.database import Database
-from macrostrat.database.transfer import pg_dump_to_file
-from macrostrat.database.transfer.utils import raw_database_url
-from macrostrat.utils import get_logger
-from macrostrat.utils.shell import run
-
-from ..database.utils import engine_for_db_name
 from .defs import (
     StatementCounter,
     apply_schema_for_environment,
@@ -35,6 +34,7 @@
 # First, register all migrations
 # NOTE: right now, this is quite implicit.
 from .migration_system import load_migrations
+from ..database.utils import engine_for_db_name
 
 log = get_logger(__name__)
 
@@ -93,8 +93,17 @@ def plan():
             f"[dim]Writing {len(m.statements)} proposed changes to [bold]{out_file}[/]"
         )
 
+        statements = [
+            format_sql(s, reindent_aligned=True, compact=True, keyword_case="upper")
+            for s in m.statements
+        ]
+
         with open(out_file, "w") as f:
-            f.write("\n".join(m.statements))
+            for statement in statements:
+                f.write(statement + "\n")
+                # Extra newline after statements that span multiple lines
+                if "\n" in statement:
+                    f.write("\n")
 
 
 @schema_app.command(rich_help_panel="Automated migrations")

py-modules/cli/macrostrat/cli/schema_management/defs.py

@@ -1,14 +1,14 @@
 from contextlib import contextmanager
 
 import docker
+from macrostrat.database import Database
+from macrostrat.dinosaur.upgrade_cluster.utils import database_cluster
 from results.dbdiff.statements import check_for_drop
 from results.schemainspect.pg import PostgreSQL
 from results.schemainspect.pg.obj import PROPS
 from rich import print
 
 from macrostrat.core.config import settings
-from macrostrat.database import Database
-from macrostrat.dinosaur.upgrade_cluster.utils import database_cluster
 
 
 def is_unsafe_statement(s: str) -> bool:
@@ -51,18 +51,26 @@ def schema_dirs_for_environment(env: str):
 
 
 def apply_schema_for_environment(
-    db: Database, env: str, *, recursive: bool = True, statement_filter=None
+    db: Database,
+    env: str,
+    *,
+    recursive: bool = True,
+    statement_filter=lambda s, p: True,
 ):
     for env_dir in schema_dirs_for_environment(env):
         schema_dir = env_dir
         if not schema_dir.exists():
             continue
-        fixtures = sorted(list(schema_dir.glob("*.sql")))
+
+        func = schema_dir.rglob if recursive else schema_dir.glob
+        fixtures = sorted(list(func("*.sql")))
         fixtures = [f for f in fixtures if not f.name.endswith(".plan.sql")]
 
         if len(fixtures) == 0:
             continue
-        db.run_fixtures(fixtures, recursive=recursive, statement_filter=filter)
+        db.run_fixtures(
+            fixtures, recursive=recursive, statement_filter=statement_filter
+        )
 
 
 @contextmanager