"""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.
"""
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
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")
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 = _format_sequential_version(parts[0].strip())
end = _format_sequential_version(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
@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).
"""
source_migrations: list[tuple[str, Path]]
target_version: str
target_path: Path
description: str
source_versions: list[str]