feat(migration): improved template customization (#245)

Improved python migration file support/templating. Allow all fields to be configurable (such as `title` and `author` of migration)

Author: cofin

docs/usage/migrations.rst

@@ -139,6 +139,81 @@ All database configs (sync and async) provide these migration methods:
 ``get_current_migration(verbose=False)``
    Get the current migration version.
 
+Template Profiles & Author Metadata
+===================================
+
+Migrations inherit their header text, metadata comments, and default file format
+from ``migration_config["templates"]``. Each project can define multiple
+profiles and select one globally:
+
+.. code-block:: python
+
+   migration_config={
+       "default_format": "py",      # CLI default when --format omitted
+       "title": "Acme Migration",    # Shared title for all templates
+       "author": "env:SQLSPEC_AUTHOR",  # Read from environment variable
+       "templates": {
+           "sql": {
+               "header": "-- {title} - {message}",
+               "metadata": ["-- Version: {version}", "-- Owner: {author}"],
+               "body": "-- custom SQL body"
+           },
+           "py": {
+               "docstring": """{title}\nDescription: {description}""",
+               "imports": ["from typing import Iterable"],
+               "body": """def up(context: object | None = None) -> str | Iterable[str]:
+    return "SELECT 1"
+
+def down(context: object | None = None) -> str | Iterable[str]:
+    return "DROP TABLE example;"
+"""
+           }
+       }
+   }
+
+Template fragments accept the following variables:
+
+- ``{title}`` – shared template title
+- ``{version}`` – generated revision identifier
+- ``{message}`` – CLI/command message
+- ``{description}`` – message fallback used in logs and docstrings
+- ``{created_at}`` – UTC timestamp in ISO 8601 format
+- ``{author}`` – resolved author string
+- ``{adapter}`` – config driver class (useful for docstrings)
+- ``{project_slug}`` / ``{slug}`` – sanitized project and message slugs
+
+Missing placeholders raise ``TemplateValidationError`` so mistakes are caught
+immediately. SQL templates list metadata rows (``metadata``) and a ``body``
+block. Python templates expose ``docstring``, optional ``imports``, and ``body``.
+
+Author attribution can be controlled via ``migration_config["author"]``:
+
+- Literal strings (``"Data Platform"``) are stamped verbatim
+- ``"env:VAR_NAME"`` pulls from the environment and fails fast if unset
+- ``"callable:pkg.module:get_author"`` invokes a helper that can inspect the
+  config or environment when determining the author string
+- ``"git"`` reads git user.name/email; ``"system"`` uses ``$USER``
+
+CLI Enhancements
+----------------
+
+``sqlspec create-migration`` (and ``litestar database create-migration``)
+accept ``--format`` / ``--file-type`` flags:
+
+.. code-block:: bash
+
+   sqlspec --config myapp.config create-migration -m "Add seed data" --format py
+
+When omitted, the CLI uses ``migration_config["default_format"]`` (``"sql"`` by default).
+Upgrade/downgrade commands now echo ``{version}: {description}``, so the rich
+description captured in templates is visible during deployments and matches the
+continue-on-error logs.
+
+The default Python template ships with both ``up`` and ``down`` functions that
+accept an optional ``context`` argument. When migrations run via SQLSpec, that
+parameter receives the active ``MigrationContext`` so you can reach the config
+or connection objects directly inside your migration logic.
+
 ``create_migration(message, file_type="sql")``
    Create a new migration file.
 
@@ -452,7 +527,8 @@ SQLSpec uses a tracking table to record applied migrations:
    renamed migrations (e.g., timestamp → sequential conversion).
 
 ``applied_by``
-   Unix username of user who applied the migration.
+   Author string recorded for the migration. Defaults to the git user/system
+   account but can be overridden via ``migration_config["author"]``.
 
 Schema Migration
 ----------------

pyproject.toml

@@ -503,6 +503,7 @@ split-on-trailing-comma = false
 "docs/examples/**" = ["T201"]
 "sqlspec/builder/mixins/**/*.*" = ["SLF001"]
 "sqlspec/extensions/adk/converters.py" = ["S403"]
+"sqlspec/migrations/utils.py" = ["S404"]
 "tests/**/*.*" = [
   "A",
   "ARG",

sqlspec/cli.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 import rich_click as click
+from click.core import ParameterSource
 
 if TYPE_CHECKING:
     from rich_click import Group
@@ -80,6 +81,16 @@ def sqlspec_group(ctx: "click.Context", config: str, validate_config: bool) -> N
     return sqlspec_group
 
 
+def _ensure_click_context() -> "click.Context":
+    """Return the active Click context, raising if missing (for type-checkers)."""
+
+    context = click.get_current_context()
+    if context is None:  # pragma: no cover - click guarantees context in commands
+        msg = "SQLSpec CLI commands require an active Click context"
+        raise RuntimeError(msg)
+    return cast("click.Context", context)
+
+
 def add_migration_commands(database_group: "Group | None" = None) -> "Group":
     """Add migration commands to the database group.
 
@@ -270,17 +281,12 @@ def show_database_revision(  # pyright: ignore[reportUnusedFunction]
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _show_current_revision() -> None:
             # Check if this is a multi-config operation
             configs_to_process = process_multiple_configs(
-                cast("click.Context", ctx),
-                bind_key,
-                include,
-                exclude,
-                dry_run=False,
-                operation_name="show current revision",
+                ctx, bind_key, include, exclude, dry_run=False, operation_name="show current revision"
             )
 
             if configs_to_process is not None:
@@ -300,7 +306,7 @@ async def _show_current_revision() -> None:
                         console.print(f"[red]✗ Failed to get current revision for {config_name}: {e}[/]")
             else:
                 console.rule("[yellow]Listing current revision[/]", align="left")
-                sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+                sqlspec_config = get_config_by_bind_key(ctx, bind_key)
                 migration_commands = create_migration_commands(config=sqlspec_config)
                 await maybe_await(migration_commands.current(verbose=verbose))
 
@@ -327,17 +333,12 @@ def downgrade_database(  # pyright: ignore[reportUnusedFunction]
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _downgrade_database() -> None:
             # Check if this is a multi-config operation
             configs_to_process = process_multiple_configs(
-                cast("click.Context", ctx),
-                bind_key,
-                include,
-                exclude,
-                dry_run=dry_run,
-                operation_name=f"downgrade to {revision}",
+                ctx, bind_key, include, exclude, dry_run=dry_run, operation_name=f"downgrade to {revision}"
             )
 
             if configs_to_process is not None:
@@ -371,7 +372,7 @@ async def _downgrade_database() -> None:
                     else Confirm.ask(f"Are you sure you want to downgrade the database to the `{revision}` revision?")
                 )
                 if input_confirmed:
-                    sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+                    sqlspec_config = get_config_by_bind_key(ctx, bind_key)
                     migration_commands = create_migration_commands(config=sqlspec_config)
                     await maybe_await(migration_commands.downgrade(revision=revision, dry_run=dry_run))
 
@@ -402,7 +403,7 @@ def upgrade_database(  # pyright: ignore[reportUnusedFunction]
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _upgrade_database() -> None:
             # Report execution mode when specified
@@ -411,7 +412,7 @@ async def _upgrade_database() -> None:
 
             # Check if this is a multi-config operation
             configs_to_process = process_multiple_configs(
-                cast("click.Context", ctx), bind_key, include, exclude, dry_run, operation_name=f"upgrade to {revision}"
+                ctx, bind_key, include, exclude, dry_run, operation_name=f"upgrade to {revision}"
             )
 
             if configs_to_process is not None:
@@ -449,7 +450,7 @@ async def _upgrade_database() -> None:
                     )
                 )
                 if input_confirmed:
-                    sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+                    sqlspec_config = get_config_by_bind_key(ctx, bind_key)
                     migration_commands = create_migration_commands(config=sqlspec_config)
                     await maybe_await(
                         migration_commands.upgrade(revision=revision, auto_sync=not no_auto_sync, dry_run=dry_run)
@@ -465,10 +466,10 @@ def stamp(bind_key: str | None, revision: str) -> None:  # pyright: ignore[repor
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _stamp() -> None:
-            sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+            sqlspec_config = get_config_by_bind_key(ctx, bind_key)
             migration_commands = create_migration_commands(config=sqlspec_config)
             await maybe_await(migration_commands.stamp(revision=revision))
 
@@ -488,7 +489,7 @@ def init_sqlspec(  # pyright: ignore[reportUnusedFunction]
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _init_sqlspec() -> None:
             console.rule("[yellow]Initializing database migrations.", align="left")
@@ -498,11 +499,7 @@ async def _init_sqlspec() -> None:
                 else Confirm.ask("[bold]Are you sure you want initialize migrations for the project?[/]")
             )
             if input_confirmed:
-                configs = (
-                    [get_config_by_bind_key(cast("click.Context", ctx), bind_key)]
-                    if bind_key is not None
-                    else cast("click.Context", ctx).obj["configs"]
-                )
+                configs = [get_config_by_bind_key(ctx, bind_key)] if bind_key is not None else ctx.obj["configs"]
 
                 for config in configs:
                     migration_config = getattr(config, "migration_config", {})
@@ -519,17 +516,25 @@ async def _init_sqlspec() -> None:
     )
     @bind_key_option
     @click.option("-m", "--message", default=None, help="Revision message")
+    @click.option(
+        "--format",
+        "--file-type",
+        "file_format",
+        type=click.Choice(["sql", "py"]),
+        default=None,
+        help="File format for the generated migration (defaults to template profile)",
+    )
     @no_prompt_option
     def create_revision(  # pyright: ignore[reportUnusedFunction]
-        bind_key: str | None, message: str | None, no_prompt: bool
+        bind_key: str | None, message: str | None, file_format: str | None, no_prompt: bool
     ) -> None:
         """Create a new database revision."""
         from rich.prompt import Prompt
 
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _create_revision() -> None:
             console.rule("[yellow]Creating new migration revision[/]", align="left")
@@ -539,9 +544,11 @@ async def _create_revision() -> None:
                     "new migration" if no_prompt else Prompt.ask("Please enter a message describing this revision")
                 )
 
-            sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+            sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+            param_source = ctx.get_parameter_source("file_format")
+            effective_format = None if param_source is ParameterSource.DEFAULT else file_format
             migration_commands = create_migration_commands(config=sqlspec_config)
-            await maybe_await(migration_commands.revision(message=message_text))
+            await maybe_await(migration_commands.revision(message=message_text, file_type=effective_format))
 
         run_(_create_revision)()
 
@@ -557,11 +564,11 @@ def fix_migrations(  # pyright: ignore[reportUnusedFunction]
         from sqlspec.migrations.commands import create_migration_commands
         from sqlspec.utils.sync_tools import run_
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         async def _fix_migrations() -> None:
             console.rule("[yellow]Migration Fix Command[/]", align="left")
-            sqlspec_config = get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+            sqlspec_config = get_config_by_bind_key(ctx, bind_key)
             migration_commands = create_migration_commands(config=sqlspec_config)
             await maybe_await(migration_commands.fix(dry_run=dry_run, update_database=not no_database, yes=yes))
 
@@ -573,20 +580,20 @@ def show_config(bind_key: str | None = None) -> None:  # pyright: ignore[reportU
         """Show and display all configurations with migrations enabled."""
         from rich.table import Table
 
-        ctx = click.get_current_context()
+        ctx = _ensure_click_context()
 
         # If bind_key is provided, filter to only that config
         if bind_key is not None:
-            get_config_by_bind_key(cast("click.Context", ctx), bind_key)
+            get_config_by_bind_key(ctx, bind_key)
             # Convert single config to list format for compatibility
-            all_configs = cast("click.Context", ctx).obj["configs"]
+            all_configs = ctx.obj["configs"]
             migration_configs = []
             for cfg in all_configs:
                 config_name = cfg.bind_key
                 if config_name == bind_key and hasattr(cfg, "migration_config") and cfg.migration_config:
                     migration_configs.append((config_name, cfg))
         else:
-            migration_configs = get_configs_with_migrations(cast("click.Context", ctx))
+            migration_configs = get_configs_with_migrations(ctx)
 
         if not migration_configs:
             console.print("[yellow]No configurations with migrations detected.[/]")