"""Migration command implementations for SQLSpec.
This module provides the main command interface for database migrations.
"""
import functools
import inspect
import logging
import time
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
from rich.console import Console
from rich.table import Table
from sqlspec.builder import sql
from sqlspec.exceptions import MigrationError
from sqlspec.migrations.base import BaseMigrationCommands
from sqlspec.migrations.context import MigrationContext
from sqlspec.migrations.fix import MigrationFixer
from sqlspec.migrations.runner import AsyncMigrationRunner, SyncMigrationRunner
from sqlspec.migrations.squash import MigrationSquasher
from sqlspec.migrations.utils import create_migration_file
from sqlspec.migrations.validation import validate_migration_order
from sqlspec.migrations.version import generate_conversion_map, generate_timestamp_version
from sqlspec.observability import resolve_db_system
from sqlspec.utils.logging import get_logger, log_with_context
if TYPE_CHECKING:
from pathlib import Path
from sqlspec.config import AsyncConfigT, SyncConfigT
from sqlspec.migrations.base import AppliedMigrationRecord, LoadedMigrationMetadata
__all__ = ("AsyncMigrationCommands", "SyncMigrationCommands", "create_migration_commands")
logger = get_logger("sqlspec.migrations.commands")
console = Console()
P = ParamSpec("P")
R = TypeVar("R")
MetadataBuilder = Callable[[dict[str, Any]], tuple[str | None, dict[str, Any]]]
def _bind_arguments(signature: inspect.Signature, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
bound = signature.bind_partial(*args, **kwargs)
arguments = dict(bound.arguments)
arguments.pop("self", None)
return arguments
def _with_command_span(
event: str, metadata_fn: "MetadataBuilder | None" = None, *, dry_run_param: str | None = "dry_run"
) -> Callable[[Callable[P, R]], Callable[P, R]]:
"""Attach span lifecycle and command metric management to command methods."""
metric_prefix = f"migrations.command.{event}"
def decorator(func: Callable[P, R]) -> Callable[P, R]:
signature = inspect.signature(func)
def _prepare(self: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[Any, bool, Any]:
runtime = self._runtime
metadata_args = _bind_arguments(signature, args, kwargs)
dry_run = False
if dry_run_param is not None:
dry_run = bool(metadata_args.get(dry_run_param, False))
metadata: dict[str, Any] | None = None
version: str | None = None
span = None
if runtime is not None:
runtime.increment_metric(f"{metric_prefix}.invocations")
if dry_run_param is not None and dry_run:
runtime.increment_metric(f"{metric_prefix}.dry_run")
if metadata_fn is not None:
version, metadata = metadata_fn(metadata_args)
span = runtime.start_migration_span(f"command.{event}", version=version, metadata=metadata)
return runtime, dry_run, span
def _finalize(
self: Any,
runtime: Any,
span: Any,
start: float,
error: "Exception | None",
recorded_error: bool,
dry_run: bool,
) -> None:
command_error = self._last_command_error
self._last_command_error = None
command_metrics = self._last_command_metrics
self._last_command_metrics = None
if runtime is None:
return
if command_error is not None and not recorded_error:
runtime.increment_metric(f"{metric_prefix}.errors")
if not dry_run and command_metrics:
for metric, value in command_metrics.items():
runtime.increment_metric(f"{metric_prefix}.{metric}", value)
duration_ms = int((time.perf_counter() - start) * 1000)
runtime.end_migration_span(span, duration_ms=duration_ms, error=error or command_error)
if inspect.iscoroutinefunction(func):
@functools.wraps(func)
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
self = args[0]
runtime, dry_run, span = _prepare(self, args, kwargs)
start = time.perf_counter()
error: Exception | None = None
error_recorded = False
try:
async_func = cast("Callable[P, Awaitable[R]]", func)
return await async_func(*args, **kwargs)
except Exception as exc: # pragma: no cover
error = exc
if runtime is not None:
runtime.increment_metric(f"{metric_prefix}.errors")
error_recorded = True
raise
finally:
_finalize(self, runtime, span, start, error, error_recorded, dry_run)
return cast("Callable[P, R]", async_wrapper)
@functools.wraps(func)
def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
self = args[0]
runtime, dry_run, span = _prepare(self, args, kwargs)
start = time.perf_counter()
error: Exception | None = None
error_recorded = False
try:
return func(*args, **kwargs)
except Exception as exc: # pragma: no cover
error = exc
if runtime is not None:
runtime.increment_metric(f"{metric_prefix}.errors")
error_recorded = True
raise
finally:
_finalize(self, runtime, span, start, error, error_recorded, dry_run)
return cast("Callable[P, R]", sync_wrapper)
return decorator
def _command_metadata(args: dict[str, Any]) -> tuple[str | None, dict[str, Any]]:
"""Build span metadata (revision, dry_run) shared by upgrade and downgrade commands."""
revision = cast("str | None", args.get("revision"))
metadata = {"dry_run": str(args.get("dry_run", False)).lower()}
return revision, metadata
class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
"""Synchronous migration commands."""
def __init__(self, config: "SyncConfigT") -> None:
"""Initialize migration commands.
Args:
config: The SQLSpec configuration.
"""
super().__init__(config)
self.tracker = self._create_tracker()
# Create context with extension configurations
context = MigrationContext.from_config(config)
context.extension_config = self.extension_configs
self.runner = SyncMigrationRunner(
self.migrations_path,
self._discover_extension_migrations(),
context,
self.extension_configs,
runtime=self._runtime,
description_hints=self._template_settings.description_hints,
)