"""Migration execution engine for SQLSpec."""
import ast
import hashlib
import logging
import re
import time
from abc import abstractmethod
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, cast, overload
from sqlspec.core import SQL
from sqlspec.loader import SQLFileLoader
from sqlspec.migrations.context import MigrationContext
from sqlspec.migrations.loaders import _load_migration_sql, get_migration_loader
from sqlspec.migrations.templates import TemplateDescriptionHints
from sqlspec.migrations.utils import resolve_default_schema as _resolve_default_schema
from sqlspec.migrations.version import _format_sequential_version, parse_extension_stem, parse_version
from sqlspec.observability import resolve_db_system
from sqlspec.utils.logging import get_logger, log_with_context
from sqlspec.utils.sync_tools import async_
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from sqlspec.config import DatabaseConfigProtocol
from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
from sqlspec.migrations.base import LoadedMigrationMetadata
from sqlspec.observability import ObservabilityRuntime
__all__ = ("AsyncMigrationRunner", "SyncMigrationRunner", "create_migration_runner")
logger = get_logger("sqlspec.migrations.runner")
class _CachedMigrationMetadata:
"""Cached migration metadata keyed by file path."""
__slots__ = ("metadata", "mtime_ns", "size")
def __init__(self, metadata: "dict[str, Any]", mtime_ns: int, size: int) -> None:
self.metadata = metadata
self.mtime_ns = mtime_ns
self.size = size
def clone(self) -> "dict[str, Any]":
return dict(self.metadata)
class _MigrationFileEntry:
"""Represents a migration file discovered during directory scanning."""
__slots__ = ("extension_name", "path")
def __init__(self, path: Path, extension_name: "str | None") -> None:
self.path = path
self.extension_name = extension_name
class BaseMigrationRunner:
"""Base migration runner with common functionality shared between sync and async implementations."""
def __init__(
self,
migrations_path: Path,
extension_migrations: "dict[str, Path] | None" = None,
context: "MigrationContext | None" = None,
extension_configs: "dict[str, dict[str, Any]] | None" = None,
runtime: "ObservabilityRuntime | None" = None,
description_hints: "TemplateDescriptionHints | None" = None,
summary_only: bool = False,
use_logger: bool = False,
) -> None:
"""Initialize the migration runner.
Args:
migrations_path: Path to the directory containing migration files.
extension_migrations: Optional mapping of extension names to their migration paths.
context: Optional migration context for Python migrations.
extension_configs: Optional mapping of extension names to their configurations.
runtime: Observability runtime shared with command/context consumers.
description_hints: Hints for extracting migration descriptions.
summary_only: Whether summary-only logging is enabled.
use_logger: Whether to emit log output. Defaults to False (CLI mode).
"""
self.migrations_path = migrations_path
self.extension_migrations = extension_migrations or {}
self.runtime = runtime
self.loader = SQLFileLoader(runtime=runtime)
self._extension_sql_loaders: dict[str, SQLFileLoader] = {}
self.project_root: Path | None = None
self.context = context
self.extension_configs = extension_configs or {}
self._listing_digest: str | None = None
self._listing_cache: list[tuple[str, Path]] | None = None
self._listing_signatures: dict[str, tuple[int, int]] = {}
self._metadata_cache: dict[str, _CachedMigrationMetadata] = {}
self.description_hints = description_hints or TemplateDescriptionHints()
self.summary_only = summary_only
self.use_logger = use_logger
def set_summary_only(self, value: bool) -> None:
"""Set summary-only logging behavior for migration runner."""
self.summary_only = value
def set_use_logger(self, value: bool) -> None:
"""Set whether to emit log output.
Args:
value: True to enable logging, False for silent mode (CLI default).
"""
self.use_logger = value
def set_driver(self, driver: "SyncDriverAdapterBase | AsyncDriverAdapterBase") -> None:
"""Bind the active session driver to the migration context."""
if self.context is not None:
self.context.driver = driver
def _log_migration_event(self, level: int, event: str, **extra_fields: Any) -> None:
"""Log migration events, respecting use_logger and summary_only settings."""
if not self.use_logger:
return
if self.summary_only and level == logging.INFO:
return
log_with_context(logger, level, event, **extra_fields)
def _metric(self, name: str, amount: float = 1.0) -> None:
if self.runtime is None:
return
self.runtime.increment_metric(name, amount)
def _iter_directory_entries(self, base_path: Path, extension_name: "str | None") -> "list[_MigrationFileEntry]":
"""Collect migration files discovered under a base path."""
if not base_path.exists():
return []
entries: list[_MigrationFileEntry] = []
for pattern in ("*.sql", "*.py"):
for file_path in sorted(base_path.glob(pattern)):
if file_path.name.startswith("."):
continue
entries.append(_MigrationFileEntry(path=file_path, extension_name=extension_name))
return entries
def _collect_listing_entries(self) -> "tuple[list[_MigrationFileEntry], dict[str, tuple[int, int]], str]":
"""Gather migration files, stat signatures, and digest for cache validation."""
entries: list[_MigrationFileEntry] = []
signatures: dict[str, tuple[int, int]] = {}
digest_source = hashlib.md5(usedforsecurity=False)
for entry in self._iter_directory_entries(self.migrations_path, None):
self._record_entry(entry, entries, signatures, digest_source)
for ext_name, ext_path in self.extension_migrations.items():
for entry in self._iter_directory_entries(ext_path, ext_name):
self._record_entry(entry, entries, signatures, digest_source)
return entries, signatures, digest_source.hexdigest()
def _record_entry(
self,
entry: _MigrationFileEntry,
entries: "list[_MigrationFileEntry]",
signatures: "dict[str, tuple[int, int]]",
digest_source: Any,
) -> None:
"""Record entry metadata for cache decisions."""
try:
stat_result = entry.path.stat()
except FileNotFoundError:
return
path_str = str(entry.path)
token = (stat_result.st_mtime_ns, stat_result.st_size)
signatures[path_str] = token
digest_source.update(path_str.encode("utf-8"))
digest_source.update(f"{token[0]}:{token[1]}".encode())
entries.append(entry)
def _build_sorted_listing(self, entries: "list[_MigrationFileEntry]") -> "list[tuple[str, Path]]":
"""Construct sorted migration listing from directory entries."""
migrations: list[tuple[str, Path]] = []
for entry in entries:
version = self._extract_version(entry.path.name)
if not version:
continue
if entry.extension_name:
version = f"ext_{entry.extension_name}_{version}"
migrations.append((version, entry.path))
def version_sort_key(migration_tuple: "tuple[str, Path]") -> "Any":
version_str = migration_tuple[0]
try:
return parse_version(version_str)
except ValueError:
return version_str
return sorted(migrations, key=version_sort_key)
def _log_listing_invalidation(
self, previous: "dict[str, tuple[int, int]]", current: "dict[str, tuple[int, int]]"
) -> None:
"""Log cache invalidation details at INFO level."""
prev_keys = set(previous)
curr_keys = set(current)
added = curr_keys - prev_keys
removed = prev_keys - curr_keys
modified = {key for key in prev_keys & curr_keys if previous[key] != current[key]}
self._log_migration_event(
logging.INFO,
"migration.listing.invalidated",
added_count=len(added),
removed_count=len(removed),
modified_count=len(modified),
)
self._metric("migrations.listing.cache_invalidations")
if added:
self._metric("migrations.listing.added", float(len(added)))
if removed:
self._metric("migrations.listing.removed", float(len(removed)))
if modified:
self._metric("migrations.listing.modified", float(len(modified)))
def _extract_version(self, filename: str) -> "str | None":
"""Extract version from filename.
Supports sequential (0001), timestamp (20251011120000), and extension-prefixed
(ext_litestar_0001) version formats.
Args:
filename: The migration filename.
Returns:
The extracted version string or None.
"""
timestamp_min_length = 4
name_without_ext = filename.rsplit(".", 1)[0]
if name_without_ext.startswith("ext_"):
extension = parse_extension_stem(name_without_ext)
if extension is None:
return None
ext_name, ext_version = extension
return f"ext_{ext_name}_{ext_version}"
parts = name_without_ext.split("_", 1)
if parts and parts[0].isdigit():
return parts[0] if len(parts[0]) > timestamp_min_length else _format_sequential_version(parts[0])
return None
def calculate_checksum(self, content: str) -> str:
"""Calculate MD5 checksum of migration content.
Canonicalizes content by excluding query name headers that change during
fix command (migrate-{version}-up/down). This ensures checksums remain
stable when converting timestamp versions to sequential format.
Args:
content: The migration file content.
Returns:
MD5 checksum hex string.
"""
canonical_content = re.sub(r"^--\s*name:\s*migrate-[^-]+-(?:up|down)\s*$", "", content, flags=re.MULTILINE)
return hashlib.md5(canonical_content.encode()).hexdigest() # noqa: S324
@abstractmethod
def load_migration(self, file_path: Path) -> "LoadedMigrationMetadata | Awaitable[LoadedMigrationMetadata]":
"""Load a migration file and extract its components.
Args:
file_path: Path to the migration file.
Returns:
Dictionary containing migration metadata and queries.
For async implementations, returns a coroutine.
"""
def _load_migration_listing(self) -> "list[tuple[str, Path]]":
"""Build the cached migration listing shared by sync/async runners."""
entries, signatures, digest = self._collect_listing_entries()
cached_listing = self._listing_cache
if cached_listing is not None and self._listing_digest == digest:
self._metric("migrations.listing.cache_hit")
self._metric("migrations.listing.files_cached", float(len(cached_listing)))
self._log_migration_event(logging.DEBUG, "migration.listing.cache_hit", file_count=len(cached_listing))
return cached_listing
files = self._build_sorted_listing(entries)
previous_digest = self._listing_digest
previous_signatures = self._listing_signatures
self._metric("migrations.listing.cache_miss")
self._metric("migrations.listing.files_scanned", float(len(files)))
self._listing_cache = files
self._listing_signatures = signatures
self._listing_digest = digest
if previous_digest is None:
self._log_migration_event(logging.DEBUG, "migration.listing.cache_primed", file_count=len(files))
else:
self._log_listing_invalidation(previous_signatures, signatures)
return files
@abstractmethod
def get_migration_files(self) -> "list[tuple[str, Path]] | Awaitable[list[tuple[str, Path]]]":
"""Get all migration files sorted by version."""
def _load_metadata(self, file_path: Path, version: "str | None" = None) -> "LoadedMigrationMetadata":
"""Load common migration metadata that doesn't require async operations.
Args:
file_path: Path to the migration file.
version: Optional pre-extracted version (preserves prefixes like ext_adk_0001).
Returns:
Partial migration metadata dictionary.
"""
cache_key = str(file_path)
stat_result = file_path.stat()
cached_metadata = self._metadata_cache.get(cache_key)
if (
cached_metadata
and cached_metadata.mtime_ns == stat_result.st_mtime_ns
and cached_metadata.size == stat_result.st_size
):
self._metric("migrations.metadata.cache_hit")
self._log_migration_event(logging.DEBUG, "migration.metadata.cache_hit", file_path=cache_key)
metadata = cast("LoadedMigrationMetadata", cached_metadata.clone())
metadata["file_path"] = file_path
return metadata
self._metric("migrations.metadata.cache_miss")
self._metric("migrations.metadata.bytes", float(stat_result.st_size))
content = file_path.read_text(encoding="utf-8")
checksum = self.calculate_checksum(content)
if version is None:
version = self._extract_version(file_path.name)
description = self._extract_description(content, file_path)
if not description:
description = file_path.stem.split("_", 1)[1] if "_" in file_path.stem else ""
transactional_match = re.search(
r"^--\s*transactional:\s*(true|false)\s*$", content, re.MULTILINE | re.IGNORECASE
)
transactional = None
if transactional_match:
transactional = transactional_match.group(1).lower() == "true"
metadata = cast(
"LoadedMigrationMetadata",
{
"version": version,
"description": description,
"file_path": file_path,
"checksum": checksum,
"content": content,
"transactional": transactional,
},
)
self._metadata_cache[cache_key] = _CachedMigrationMetadata(
metadata=dict(metadata), mtime_ns=stat_result.st_mtime_ns, size=stat_result.st_size
)
if cached_metadata:
self._log_migration_event(logging.DEBUG, "migration.metadata.cache_invalidated", file_path=cache_key)
else:
self._log_migration_event(logging.DEBUG, "migration.metadata.cached", file_path=cache_key)
return metadata
def _extract_description(self, content: str, file_path: Path) -> str:
if file_path.suffix == ".sql":
return self._extract_sql_description(content)
if file_path.suffix == ".py":
return self._extract_python_description(content)
return ""
def _extract_sql_description(self, content: str) -> str:
keys = self.description_hints.sql_keys
for line in content.splitlines():
stripped = line.strip()
if not stripped:
continue
if stripped.startswith("--"):
body = stripped.lstrip("-").strip()
if not body:
continue
if ":" in body:
key, value = body.split(":", 1)
if key.strip() in keys:
return value.strip()
continue
break
return ""
def _extract_python_description(self, content: str) -> str:
try:
module = ast.parse(content)
except SyntaxError:
return ""
docstring = ast.get_docstring(module) or ""
keys = self.description_hints.python_keys
for line in docstring.splitlines():
stripped = line.strip()
if not stripped:
continue
if ":" in stripped:
key, value = stripped.split(":", 1)
if key.strip() in keys:
return value.strip()
return stripped
return ""
def _migration_context(self, file_path: Path) -> "MigrationContext | None":
"""Get the appropriate context for a migration file.
Args:
file_path: Path to the migration file.
Returns:
Migration context to use, or None to use default.
"""
context_to_use = self.context
if context_to_use and file_path.name.startswith("ext_"):
version = self._extract_version(file_path.name)
if version and version.startswith("ext_"):
extension = parse_extension_stem(version)
if extension is not None:
ext_name = extension[0]
if ext_name in self.extension_configs:
context_to_use = MigrationContext(
dialect=self.context.dialect if self.context else None,
config=self.context.config if self.context else None,
driver=self.context.driver if self.context else None,
metadata=self.context.metadata.copy() if self.context and self.context.metadata else {},
extension_config=self.extension_configs[ext_name],
)
for ext_name, ext_path in self.extension_migrations.items():
if file_path.parent == ext_path:
if ext_name in self.extension_configs and self.context:
context_to_use = MigrationContext(
config=self.context.config,
dialect=self.context.dialect,
driver=self.context.driver,
metadata=self.context.metadata.copy() if self.context.metadata else {},
extension_config=self.extension_configs[ext_name],
)
break
return context_to_use
def should_use_transaction(
self, migration: "LoadedMigrationMetadata", config: "DatabaseConfigProtocol[Any, Any, Any]"
) -> bool:
"""Determine if migration should run in a transaction.
Args:
migration: Migration metadata dictionary.
config: The database configuration instance.
Returns:
True if migration should be wrapped in a transaction.
"""
if not config.supports_transactional_ddl:
return False
if migration.get("transactional") is not None:
return bool(migration["transactional"])
migration_config = cast("dict[str, Any]", config.migration_config) or {}
return bool(migration_config.get("transactional", True))
def _resolve_default_schema(self) -> str | None:
"""Return the configured default schema for migration execution."""
config = self.context.config if self.context else None
migration_config = cast("dict[str, Any] | None", getattr(config, "migration_config", None))
return _resolve_default_schema(migration_config)
def _resolve_use_transaction(self, migration: "LoadedMigrationMetadata", use_transaction: "bool | None") -> bool:
"""Resolve the effective transaction flag for a migration."""
if use_transaction is None:
config = self.context.config if self.context else None
use_transaction = self.should_use_transaction(migration, config) if config else False
return bool(use_transaction)
def _log_migration_missing(
self, driver: Any, migration: "LoadedMigrationMetadata", operation: str, event: str
) -> None:
"""Record metrics and a warning for a migration with no SQL to apply."""
self._metric(f"migrations.{operation}.skipped")
self._log_migration_event(
logging.WARNING,
event,
db_system=resolve_db_system(type(driver).__name__),
version=migration.get("version"),
status="missing",
)
def _begin_migration_span(
self, driver: Any, migration: "LoadedMigrationMetadata", operation: str, event: str, use_transaction: bool
) -> Any:
"""Start the observability span for a migration and log the start event."""
runtime = self.runtime
span = None
if runtime is not None:
span = runtime.start_migration_span(operation, version=migration.get("version"))
runtime.increment_metric(f"migrations.{operation}.invocations")
self._log_migration_event(
logging.INFO,
event,
db_system=resolve_db_system(type(driver).__name__),
version=migration.get("version"),
use_transaction=use_transaction,
status="start",
)
return span
def _finish_migration_span_success(
self,
span: Any,
driver: Any,
migration: "LoadedMigrationMetadata",
operation: str,
event: str,
execution_time: int,
) -> None:
"""Close the observability span and log completion for a successful migration."""
runtime = self.runtime
if runtime is not None:
runtime.increment_metric(f"migrations.{operation}.applied")
runtime.increment_metric(f"migrations.{operation}.duration_ms", float(execution_time))
runtime.end_migration_span(span, duration_ms=execution_time)
self._log_migration_event(
logging.INFO,
event,
db_system=resolve_db_system(type(driver).__name__),
version=migration.get("version"),
duration_ms=execution_time,
status="complete",
)
def _finish_migration_span_error(
self,
span: Any,
driver: Any,
migration: "LoadedMigrationMetadata",
operation: str,
event: str,
start_time: float,
exc: Exception,
) -> None:
"""Close the observability span and log failure for a migration."""
duration_ms = int((time.perf_counter() - start_time) * 1000)
runtime = self.runtime
if runtime is not None:
runtime.increment_metric(f"migrations.{operation}.errors")
runtime.end_migration_span(span, duration_ms=duration_ms, error=exc)
self._log_migration_event(
logging.ERROR,
event,
db_system=resolve_db_system(type(driver).__name__),
version=migration.get("version"),
duration_ms=duration_ms,
error_type=type(exc).__name__,
status="failed",
)
def _check_missing_direction(self, migration: "LoadedMigrationMetadata", direction: str) -> bool:
"""Return True when a migration lacks the requested direction's query.
Called during migration loading before the ``has_*grade`` fields exist, so
an absent field is treated as "not yet known" and does not short-circuit.
Raises:
ValueError: When an upgrade query is required but absent.
"""
if f"has_{direction}grade" in migration and not migration.get(f"has_{direction}grade"):
if direction == "down":
self._log_migration_event(
logging.WARNING, "migration.downgrade.missing", version=migration.get("version")
)
return True
msg = f"Migration {migration.get('version')} has no upgrade query"
raise ValueError(msg)
return False
def _handle_migration_sql_error(
self, migration: "LoadedMigrationMetadata", direction: str, error: Exception
) -> None:
"""Swallow a downgrade loader failure with a warning; re-raise for upgrades.
Raises:
ValueError: When an upgrade query fails to load.
"""
if direction == "down":
self._log_migration_event(
logging.WARNING, "migration.downgrade.load_failed", version=migration.get("version"), error=str(error)
)
return
msg = f"Failed to load upgrade for migration {migration.get('version')}: {error}"
raise ValueError(msg) from error
def _finalize_migration_sql(self, sql_statements: Any) -> "list[str] | None":
"""Normalize loader output into a SQL statement list or None."""
if sql_statements:
return cast("list[str]", sql_statements)
return None
def _migration_sql_loader(self, file_path: Path, version: "str | None") -> SQLFileLoader:
"""Return the isolated core SQL loader for an extension migration."""
extension = parse_extension_stem(version) if version else None
if file_path.suffix != ".sql" or extension is None:
return self.loader
extension_name, _ = extension
if self.extension_migrations.get(extension_name) != file_path.parent:
return self.loader
loader = self._extension_sql_loaders.get(extension_name)
if loader is None:
loader = SQLFileLoader(runtime=self.runtime)
self._extension_sql_loaders[extension_name] = loader
return loader
class SyncMigrationRunner(BaseMigrationRunner):
"""Synchronous migration runner with pure sync methods."""
def get_migration_files(self) -> "list[tuple[str, Path]]":
"""Get all migration files sorted by version.
Returns:
List of (version, path) tuples sorted by version.
"""
return self._load_migration_listing()