sqlspec 0.26.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/migrations/context.py

@@ -3,7 +3,7 @@
 import asyncio
 import inspect
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 from sqlspec.utils.logging import get_logger
 
@@ -23,16 +23,16 @@ class MigrationContext:
     to migration functions, allowing them to generate dialect-specific SQL.
     """
 
-    config: "Optional[Any]" = None
+    config: "Any | None" = None
     """Database configuration object."""
-    dialect: "Optional[str]" = None
+    dialect: "str | None" = None
     """Database dialect (e.g., 'postgres', 'mysql', 'sqlite')."""
-    metadata: "Optional[dict[str, Any]]" = None
+    metadata: "dict[str, Any] | None" = None
     """Additional metadata for the migration."""
-    extension_config: "Optional[dict[str, Any]]" = None
+    extension_config: "dict[str, Any] | None" = None
     """Extension-specific configuration options."""
 
-    driver: "Optional[Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]]" = None
+    driver: "SyncDriverAdapterBase | AsyncDriverAdapterBase | None" = None
     """Database driver instance (available during execution)."""
 
     _execution_metadata: "dict[str, Any]" = field(default_factory=dict)
@@ -129,9 +129,6 @@ class MigrationContext:
 
         Args:
             migration_func: The migration function to validate.
-
-        Raises:
-            RuntimeError: If async function is used inappropriately.
         """
         if inspect.iscoroutinefunction(migration_func) and not self.is_async_execution and not self.is_async_driver:
             msg = (

sqlspec/migrations/fix.py

@@ -0,0 +1,199 @@
+"""Migration file fix operations for converting timestamp to sequential versions.
+
+This module provides utilities to convert timestamp-format migration files to
+sequential format, supporting the hybrid versioning workflow where development
+uses timestamps and production uses sequential numbers.
+"""
+
+import logging
+import re
+import shutil
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+__all__ = ("MigrationFixer", "MigrationRename")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MigrationRename:
+    """Represents a planned migration file rename operation.
+
+    Attributes:
+        old_path: Current file path.
+        new_path: Target file path after rename.
+        old_version: Current version string.
+        new_version: Target version string.
+        needs_content_update: Whether file content needs updating.
+            True for SQL files that contain query names.
+    """
+
+    old_path: Path
+    new_path: Path
+    old_version: str
+    new_version: str
+    needs_content_update: bool
+
+
+class MigrationFixer:
+    """Handles atomic migration file conversion operations.
+
+    Provides backup/rollback functionality and manages conversion from
+    timestamp-based migration files to sequential format.
+    """
+
+    def __init__(self, migrations_path: Path) -> None:
+        """Initialize migration fixer.
+
+        Args:
+            migrations_path: Path to migrations directory.
+        """
+        self.migrations_path = migrations_path
+        self.backup_path: Path | None = None
+
+    def plan_renames(self, conversion_map: dict[str, str]) -> list[MigrationRename]:
+        """Plan all file rename operations from conversion map.
+
+        Scans migration directory and builds list of MigrationRename objects
+        for all files that need conversion. Validates no target collisions.
+
+        Args:
+            conversion_map: Dictionary mapping old versions to new versions.
+
+        Returns:
+            List of planned rename operations.
+
+        Raises:
+            ValueError: If target file already exists or collision detected.
+        """
+        if not conversion_map:
+            return []
+
+        renames: list[MigrationRename] = []
+
+        for old_version, new_version in conversion_map.items():
+            matching_files = list(self.migrations_path.glob(f"{old_version}_*"))
+
+            for old_path in matching_files:
+                suffix = old_path.suffix
+                description = old_path.stem.replace(f"{old_version}_", "")
+
+                new_filename = f"{new_version}_{description}{suffix}"
+                new_path = self.migrations_path / new_filename
+
+                if new_path.exists() and new_path != old_path:
+                    msg = f"Target file already exists: {new_path}"
+                    raise ValueError(msg)
+
+                needs_content_update = suffix == ".sql"
+
+                renames.append(
+                    MigrationRename(
+                        old_path=old_path,
+                        new_path=new_path,
+                        old_version=old_version,
+                        new_version=new_version,
+                        needs_content_update=needs_content_update,
+                    )
+                )
+
+        return renames
+
+    def create_backup(self) -> Path:
+        """Create timestamped backup directory with all migration files.
+
+        Returns:
+            Path to created backup directory.
+
+        """
+        timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d_%H%M%S")
+        backup_dir = self.migrations_path / f".backup_{timestamp}"
+
+        backup_dir.mkdir(parents=True, exist_ok=False)
+
+        for file_path in self.migrations_path.iterdir():
+            if file_path.is_file() and not file_path.name.startswith("."):
+                shutil.copy2(file_path, backup_dir / file_path.name)
+
+        self.backup_path = backup_dir
+        return backup_dir
+
+    def apply_renames(self, renames: "list[MigrationRename]", dry_run: bool = False) -> None:
+        """Execute planned rename operations.
+
+        Args:
+            renames: List of planned rename operations.
+            dry_run: If True, log operations without executing.
+
+        """
+        if not renames:
+            return
+
+        for rename in renames:
+            if dry_run:
+                continue
+
+            if rename.needs_content_update:
+                self.update_file_content(rename.old_path, rename.old_version, rename.new_version)
+
+            rename.old_path.rename(rename.new_path)
+
+    def update_file_content(self, file_path: Path, old_version: str, new_version: str) -> None:
+        """Update SQL query names and version comments in file content.
+
+        Transforms query names and version metadata from old version to new version:
+            -- name: migrate-{old_version}-up  →  -- name: migrate-{new_version}-up
+            -- name: migrate-{old_version}-down  →  -- name: migrate-{new_version}-down
+            -- Version: {old_version}  →  -- Version: {new_version}
+
+        Creates version-specific regex patterns to avoid unintended replacements
+        of other migrate-* patterns in the file.
+
+        Args:
+            file_path: Path to file to update.
+            old_version: Old version string.
+            new_version: New version string.
+
+        """
+        content = file_path.read_text(encoding="utf-8")
+
+        up_pattern = re.compile(rf"(-- name:\s+migrate-){re.escape(old_version)}(-up)")
+        down_pattern = re.compile(rf"(-- name:\s+migrate-){re.escape(old_version)}(-down)")
+        version_pattern = re.compile(rf"(-- Version:\s+){re.escape(old_version)}")
+
+        content = up_pattern.sub(rf"\g<1>{new_version}\g<2>", content)
+        content = down_pattern.sub(rf"\g<1>{new_version}\g<2>", content)
+        content = version_pattern.sub(rf"\g<1>{new_version}", content)
+
+        file_path.write_text(content, encoding="utf-8")
+        logger.debug("Updated content in %s", file_path.name)
+
+    def rollback(self) -> None:
+        """Restore migration files from backup.
+
+        Deletes current migration files and restores from backup directory.
+        Only restores if backup exists.
+        """
+        if not self.backup_path or not self.backup_path.exists():
+            return
+
+        for file_path in self.migrations_path.iterdir():
+            if file_path.is_file() and not file_path.name.startswith("."):
+                file_path.unlink()
+
+        for backup_file in self.backup_path.iterdir():
+            if backup_file.is_file():
+                shutil.copy2(backup_file, self.migrations_path / backup_file.name)
+
+    def cleanup(self) -> None:
+        """Remove backup directory after successful conversion.
+
+        Only removes backup if it exists. Logs warning if no backup found.
+        """
+        if not self.backup_path or not self.backup_path.exists():
+            return
+
+        shutil.rmtree(self.backup_path)
+        self.backup_path = None

sqlspec/migrations/loaders.py

@@ -10,7 +10,7 @@ import types
 from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Final, Optional
+from typing import Any, Final
 
 from sqlspec.loader import SQLFileLoader as CoreSQLFileLoader
 
@@ -77,13 +77,22 @@ class SQLFileLoader(BaseMigrationLoader):
 
     __slots__ = ("sql_loader",)
 
-    def __init__(self) -> None:
-        """Initialize SQL file loader."""
-        self.sql_loader: CoreSQLFileLoader = CoreSQLFileLoader()
+    def __init__(self, sql_loader: "CoreSQLFileLoader | None" = None) -> None:
+        """Initialize SQL file loader.
+
+        Args:
+            sql_loader: Optional shared SQLFileLoader instance to reuse.
+                If not provided, creates a new instance.
+        """
+        self.sql_loader: CoreSQLFileLoader = sql_loader if sql_loader is not None else CoreSQLFileLoader()
 
     async def get_up_sql(self, path: Path) -> list[str]:
         """Extract the 'up' SQL from a SQL migration file.
 
+        The SQL file must already be loaded via validate_migration_file()
+        before calling this method. This design ensures the file is loaded
+        exactly once during the migration process.
+
         Args:
             path: Path to SQL migration file.
 
@@ -93,9 +102,6 @@ class SQLFileLoader(BaseMigrationLoader):
         Raises:
             MigrationLoadError: If migration file is invalid or missing up query.
         """
-        self.sql_loader.clear_cache()
-        self.sql_loader.load_sql(path)
-
         version = self._extract_version(path.name)
         up_query = f"migrate-{version}-up"
 
@@ -109,15 +115,16 @@ class SQLFileLoader(BaseMigrationLoader):
     async def get_down_sql(self, path: Path) -> list[str]:
         """Extract the 'down' SQL from a SQL migration file.
 
+        The SQL file must already be loaded via validate_migration_file()
+        before calling this method. This design ensures the file is loaded
+        exactly once during the migration process.
+
         Args:
             path: Path to SQL migration file.
 
         Returns:
             List containing single SQL statement for downgrade, or empty list.
         """
-        self.sql_loader.clear_cache()
-        self.sql_loader.load_sql(path)
-
         version = self._extract_version(path.name)
         down_query = f"migrate-{version}-down"
 
@@ -141,7 +148,6 @@ class SQLFileLoader(BaseMigrationLoader):
             msg = f"Invalid migration filename: {path.name}"
             raise MigrationLoadError(msg)
 
-        self.sql_loader.clear_cache()
         self.sql_loader.load_sql(path)
         up_query = f"migrate-{version}-up"
         if not self.sql_loader.has_query(up_query):
@@ -151,14 +157,31 @@ class SQLFileLoader(BaseMigrationLoader):
     def _extract_version(self, filename: str) -> str:
         """Extract version from filename.
 
+        Supports sequential (0001), timestamp (20251011120000), and extension-prefixed
+        (ext_litestar_0001) version formats.
+
         Args:
             filename: Migration filename to parse.
 
         Returns:
-            Zero-padded version string or empty string if invalid.
+            Version string or empty string if invalid.
         """
-        parts = filename.split("_", 1)
-        return parts[0].zfill(4) if parts and parts[0].isdigit() else ""
+        extension_version_parts = 3
+        timestamp_min_length = 4
+
+        name_without_ext = filename.rsplit(".", 1)[0]
+
+        if name_without_ext.startswith("ext_"):
+            parts = name_without_ext.split("_", 3)
+            if len(parts) >= extension_version_parts:
+                return f"{parts[0]}_{parts[1]}_{parts[2]}"
+            return ""
+
+        parts = name_without_ext.split("_", 1)
+        if parts and parts[0].isdigit():
+            return parts[0] if len(parts[0]) > timestamp_min_length else parts[0].zfill(4)
+
+        return ""
 
 
 class PythonFileLoader(BaseMigrationLoader):
@@ -166,9 +189,7 @@ class PythonFileLoader(BaseMigrationLoader):
 
     __slots__ = ("context", "migrations_dir", "project_root")
 
-    def __init__(
-        self, migrations_dir: Path, project_root: "Optional[Path]" = None, context: "Optional[Any]" = None
-    ) -> None:
+    def __init__(self, migrations_dir: Path, project_root: "Path | None" = None, context: "Any | None" = None) -> None:
         """Initialize Python file loader.
 
         Args:
@@ -396,7 +417,11 @@ class PythonFileLoader(BaseMigrationLoader):
 
 
 def get_migration_loader(
-    file_path: Path, migrations_dir: Path, project_root: "Optional[Path]" = None, context: "Optional[Any]" = None
+    file_path: Path,
+    migrations_dir: Path,
+    project_root: "Path | None" = None,
+    context: "Any | None" = None,
+    sql_loader: "CoreSQLFileLoader | None" = None,
 ) -> BaseMigrationLoader:
     """Factory function to get appropriate loader for migration file.
 
@@ -405,6 +430,9 @@ def get_migration_loader(
         migrations_dir: Directory containing migration files.
         project_root: Optional project root directory for Python imports.
         context: Optional migration context to pass to Python migrations.
+        sql_loader: Optional shared SQLFileLoader instance for SQL migrations.
+            When provided, SQL files are loaded using this shared instance,
+            avoiding redundant file parsing.
 
     Returns:
         Appropriate loader instance for the file type.
@@ -417,6 +445,6 @@ def get_migration_loader(
     if suffix == ".py":
         return PythonFileLoader(migrations_dir, project_root, context)
     if suffix == ".sql":
-        return SQLFileLoader()
+        return SQLFileLoader(sql_loader)
     msg = f"Unsupported migration file type: {suffix}"
     raise MigrationLoadError(msg)