"""Migration squash engine for combining multiple migrations into a single file.
This module provides utilities to consolidate multiple sequential migrations
into a single "release" migration file, following the Django-style squash workflow.
"""
import re
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
from sqlspec.exceptions import SquashValidationError
from sqlspec.migrations._backup import create_backup, remove_backup, restore_backup
from sqlspec.migrations.loaders import _load_migration_sql
from sqlspec.migrations.validation import (
validate_extension_consistency,
validate_squash_idempotency,
validate_squash_range,
)
from sqlspec.migrations.version import _format_sequential_version, parse_extension_stem, parse_version
from sqlspec.utils.logging import get_logger
from sqlspec.utils.text import slugify
if TYPE_CHECKING:
from sqlspec.migrations.runner import SyncMigrationRunner
from sqlspec.migrations.templates import MigrationTemplateSettings
__all__ = ("MigrationSquasher", "SquashPlan", "group_migrations_by_type", "parse_version_range")
logger = get_logger("sqlspec.migrations.squash")
_TRANSACTIONAL_DIRECTIVE_RE = re.compile(r"^--\s*transactional:\s*(true|false)\s*$", re.MULTILINE | re.IGNORECASE)
_SCHEMA_DIRECTIVE_RE = re.compile(r"^--\s*schema:\s*(\S+)\s*$", re.MULTILINE | re.IGNORECASE)
def parse_version_range(range_str: str) -> tuple[str, str]:
"""Parse a version range string into (start, end) tuple.
Accepts multiple formats: ``START:END``, ``START..END``, or ``START-END``.
Args:
range_str: Version range string.
Returns:
Tuple of (start_version, end_version) zero-padded to 4 digits.
Raises:
ValueError: If the format is not recognised.
"""
for sep in (":", "..", "-"):
if sep in range_str:
parts = range_str.split(sep, 1)
start = _normalize_range_endpoint(parts[0].strip())
end = _normalize_range_endpoint(parts[1].strip())
return start, end
msg = f"Invalid VERSION_RANGE format: '{range_str}'. Use START:END, START..END, or START-END"
raise ValueError(msg)
def group_migrations_by_type(migrations: list[tuple[str, Path]]) -> list[tuple[str, list[tuple[str, Path]]]]:
"""Group consecutive migrations by file type (sql or py).
Partitions a list of migrations into groups where each group contains
consecutive migrations of the same type. This enables squashing mixed
SQL and Python migrations into separate output files.
Args:
migrations: List of (version, path) tuples to group.
Returns:
List of (type, migrations) tuples where type is "sql" or "py"
and migrations is the list of (version, path) for that group.
"""
if not migrations:
return []
groups: list[tuple[str, list[tuple[str, Path]]]] = []
current_type: str | None = None
current_group: list[tuple[str, Path]] = []
for version, path in migrations:
file_type = "py" if path.suffix == ".py" else "sql"
if file_type != current_type:
if current_group and current_type is not None:
groups.append((current_type, current_group))
current_type = file_type
current_group = [(version, path)]
else:
current_group.append((version, path))
if current_group and current_type is not None:
groups.append((current_type, current_group))
return groups
[docs]
@dataclass(slots=True)
class SquashPlan:
"""Represents a planned squash operation.
Attributes:
source_migrations: List of (version, path) tuples for migrations being squashed.
target_version: The version string for the squashed migration.
target_path: Output file path for the squashed migration.
description: Combined description for the squashed migration.
source_versions: List of version strings being replaced (for tracking table updates).
transactional: Optional transactional directive preserved from source SQL migrations.
schema: Optional schema directive preserved from source SQL migrations.
"""
source_migrations: list[tuple[str, Path]]
target_version: str
target_path: Path
description: str
source_versions: list[str]
transactional: bool | None = None
schema: str | None = None
[docs]
class MigrationSquasher:
"""Core squash engine for combining migrations.
Provides functionality to plan, validate, and execute migration squash operations.
Combines multiple sequential migrations into a single file with merged UP/DOWN SQL.
"""
__slots__ = ("backup_path", "migrations_path", "runner", "template_settings")
[docs]
def __init__(
self,
migrations_path: Path,
runner: "SyncMigrationRunner",
template_settings: "MigrationTemplateSettings | None" = None,
) -> None:
"""Initialize the migration squasher.
Args:
migrations_path: Path to the migrations directory.
runner: SyncMigrationRunner instance for loading migrations.
template_settings: Optional template settings for generating squashed file.
"""
self.migrations_path = migrations_path
self.runner = runner
self.template_settings = template_settings
self.backup_path: Path | None = None
[docs]
def plan_squash(
self,
start_version: str,
end_version: str,
description: str,
*,
allow_gaps: bool = False,
output_format: str = "sql",
) -> list[SquashPlan]:
"""Plan a squash operation for a range of migrations.
For homogeneous migrations (all SQL or all Python), returns a single plan.
For mixed SQL/Python migrations or differing execution directives, returns
multiple plans - one per consecutive compatible group of migrations.
Args:
start_version: First version in the range to squash (inclusive).
end_version: Last version in the range to squash (inclusive).
description: Description for the squashed migration file.
allow_gaps: If True, allow gaps in version sequence.
output_format: Output file format ("sql" or "py").
Returns:
List of SquashPlan objects with details of planned operations.
"""
all_migrations = self.runner.get_migration_files()
source_migrations = validate_squash_range(all_migrations, start_version, end_version, allow_gaps=allow_gaps)
safe_description = slugify(description, separator="_")[:50] or "migration"
if output_format == "py":
directive_groups: list[tuple[str, bool | None, str | None, list[tuple[str, Path]]]] = [
("py", None, None, source_migrations)
]
else:
directive_groups = _group_migrations_with_directives(source_migrations)
parsed_start = parse_version(start_version)
version_counter = parsed_start.sequence or 1
ext_prefix = f"ext_{parsed_start.extension}_" if parsed_start.extension else ""
plans: list[SquashPlan] = []
for file_type, transactional, schema, group_migrations in directive_groups:
group_versions = [v for v, _ in group_migrations]
target_version = f"{ext_prefix}{_format_sequential_version(version_counter)}"
extension = ".py" if file_type == "py" else ".sql"
target_path = self.migrations_path / f"{target_version}_{safe_description}{extension}"
plans.append(
SquashPlan(
source_migrations=group_migrations,
target_version=target_version,
target_path=target_path,
description=description,
source_versions=group_versions,
transactional=transactional,
schema=schema,
)
)
version_counter += 1
return plans
[docs]
def generate_squashed_content(self, plan: SquashPlan, up_sql: list[str], down_sql: list[str]) -> str:
"""Generate the content for a squashed migration file.
Args:
plan: The SquashPlan describing the squash operation.
up_sql: List of UP SQL statements (in execution order).
down_sql: List of DOWN SQL statements (in rollback order).
Returns:
Complete SQL file content as a string.
"""
lines: list[str] = []
title = "SQLSpec Migration"
if self.template_settings and self.template_settings.profile:
title = self.template_settings.profile.title
lines.extend((
f"-- {title}",
f"-- Version: {plan.target_version}",
f"-- Description: {plan.description}",
f"-- Squashed from: {', '.join(plan.source_versions)}",
))
if plan.transactional is not None:
lines.append(f"-- transactional: {'true' if plan.transactional else 'false'}")
if plan.schema is not None:
lines.append(f"-- schema: {plan.schema}")
lines.extend(("", f"-- name: migrate-{plan.target_version}-up"))
lines.extend(statement.rstrip() for statement in up_sql)
lines.append("")
if down_sql:
lines.append(f"-- name: migrate-{plan.target_version}-down")
lines.extend(statement.rstrip() for statement in down_sql)
lines.append("")
return "\n".join(lines)
[docs]
def generate_python_squash(self, plan: SquashPlan, up_sql: list[str], down_sql: list[str]) -> str:
"""Generate Python migration file content instead of SQL.
Creates a Python migration file with up() and down() functions
that return the SQL statements as lists.
Args:
plan: The SquashPlan describing the squash operation.
up_sql: List of UP SQL statements (in execution order).
down_sql: List of DOWN SQL statements (in rollback order).
Returns:
Complete Python file content as a string.
"""
lines: list[str] = []
title = "SQLSpec Migration"
if self.template_settings and self.template_settings.profile:
title = self.template_settings.profile.title
lines.extend([
'"""' + title + ".",
"",
f"Version: {plan.target_version}",
f"Description: {plan.description}",
f"Squashed from: {', '.join(plan.source_versions)}",
'"""',
"",
])
lines.extend(["def up() -> list[str]:", ' """Return UP migration SQL statements."""', " return ["])
lines.extend(f" {statement!r}," for statement in up_sql)
lines.extend([" ]", ""])
lines.extend(["def down() -> list[str] | None:", ' """Return DOWN migration SQL statements."""'])
if down_sql:
lines.append(" return [")
lines.extend(f" {statement!r}," for statement in down_sql)
lines.extend([" ]", ""])
else:
lines.extend([" return None", ""])
return "\n".join(lines)
[docs]
def apply_squash(self, plans: list[SquashPlan], *, dry_run: bool = False, cleanup_backup: bool = True) -> None:
"""Apply the squash operation for one or more plans.
Creates backup, writes squashed files, deletes source migrations,
and cleans up backup on success. Rolls back on error.
Args:
plans: List of SquashPlan objects to execute.
dry_run: If True, no files are modified (preview only).
cleanup_backup: If True (default), remove the backup directory before returning.
Set False when callers need to keep the backup until database tracking is updated.
"""
if dry_run:
logger.debug("Dry run mode - no changes will be made")
return
ready_plans: list[SquashPlan] = []
for plan in plans:
validate_extension_consistency(plan.source_migrations)
source_files = [path for _, path in plan.source_migrations]
status = validate_squash_idempotency(source_files, plan.target_path)
if status == "already_squashed":
logger.debug("Squash target already applied: %s", plan.target_path)
continue
if status == "partial":
msg = f"Squash plan for {plan.target_path.name} is partially applied"
raise SquashValidationError(msg)
ready_plans.append(plan)
if not ready_plans:
return
self._create_backup()
try:
planned_contents: list[tuple[SquashPlan, str]] = []
for plan in ready_plans:
up_sql, down_sql = self.extract_sql(plan.source_migrations)
if plan.target_path.suffix == ".py":
content = self.generate_python_squash(plan, up_sql, down_sql)
else:
content = self.generate_squashed_content(plan, up_sql, down_sql)
planned_contents.append((plan, content))
for plan, content in planned_contents:
plan.target_path.write_text(content, encoding="utf-8")
logger.debug("Wrote squashed migration to %s", plan.target_path)
target_paths = {plan.target_path for plan in ready_plans}
all_source_paths = {
source_path
for plan in ready_plans
for _, source_path in plan.source_migrations
if source_path not in target_paths
}
for source_path in all_source_paths:
if source_path.exists():
source_path.unlink()
logger.debug("Deleted source migration %s", source_path)
if cleanup_backup:
self._cleanup_backup()
except Exception:
self._rollback_backup()
raise
def _create_backup(self) -> Path:
"""Create timestamped backup directory with all migration files.
Returns:
Path to created backup directory.
"""
self.backup_path = create_backup(self.migrations_path)
logger.debug("Created backup at %s", self.backup_path)
return self.backup_path
def _cleanup_backup(self) -> None:
"""Remove backup directory after successful operation."""
if not self.backup_path or not self.backup_path.exists():
return
remove_backup(self.backup_path)
logger.debug("Cleaned up backup at %s", self.backup_path)
self.backup_path = None
def _rollback_backup(self) -> None:
"""Restore migration files from backup on error."""
if not self.backup_path or not self.backup_path.exists():
return
backup_dir = self.backup_path
restore_backup(self.migrations_path, backup_dir, delete_backup=True)
self.backup_path = None
logger.debug("Rolled back from backup at %s", backup_dir)
def _normalize_range_endpoint(endpoint: str) -> str:
"""Normalize a squash range endpoint, preserving any extension prefix."""
if endpoint.startswith("ext_"):
parsed = parse_extension_stem(endpoint)
if parsed is not None:
ext_name, ext_version = parsed
return f"ext_{ext_name}_{_format_sequential_version(ext_version)}"
return _format_sequential_version(endpoint)
def _extract_sql_directives(path: Path) -> tuple[bool | None, str | None]:
"""Extract transactional and schema directives from a SQL migration file if it exists."""
if path.suffix != ".sql" or not path.exists():
return None, None
try:
content = path.read_text(encoding="utf-8")
except OSError:
return None, None
tx_match = _TRANSACTIONAL_DIRECTIVE_RE.search(content)
transactional = tx_match.group(1).lower() == "true" if tx_match else None
schema_match = _SCHEMA_DIRECTIVE_RE.search(content)
schema = schema_match.group(1) if schema_match else None
return transactional, schema
def _group_migrations_with_directives(
migrations: list[tuple[str, Path]],
) -> list[tuple[str, bool | None, str | None, list[tuple[str, Path]]]]:
"""Group consecutive migrations by file type and SQL execution directives."""
if not migrations:
return []
groups: list[tuple[str, bool | None, str | None, list[tuple[str, Path]]]] = []
current_key: tuple[str, bool | None, str | None] | None = None
current_group: list[tuple[str, Path]] = []
for version, path in migrations:
file_type = "py" if path.suffix == ".py" else "sql"
transactional, schema = _extract_sql_directives(path) if file_type == "sql" else (None, None)
key = (file_type, transactional, schema)
if key != current_key:
if current_group and current_key is not None:
groups.append((current_key[0], current_key[1], current_key[2], current_group))
current_key = key
current_group = [(version, path)]
else:
current_group.append((version, path))
if current_group and current_key is not None:
groups.append((current_key[0], current_key[1], current_key[2], current_group))
return groups
