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

sqlspec/migrations/context.py

@@ -0,0 +1,142 @@
+"""Migration context for passing runtime information to migrations."""
+
+import asyncio
+import inspect
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
+
+logger = get_logger("migrations.context")
+
+__all__ = ("MigrationContext",)
+
+
+@dataclass
+class MigrationContext:
+    """Context object passed to migration functions.
+
+    Provides runtime information about the database environment
+    to migration functions, allowing them to generate dialect-specific SQL.
+    """
+
+    config: "Any | None" = None
+    """Database configuration object."""
+    dialect: "str | None" = None
+    """Database dialect (e.g., 'postgres', 'mysql', 'sqlite')."""
+    metadata: "dict[str, Any] | None" = None
+    """Additional metadata for the migration."""
+    extension_config: "dict[str, Any] | None" = None
+    """Extension-specific configuration options."""
+
+    driver: "SyncDriverAdapterBase | AsyncDriverAdapterBase | None" = None
+    """Database driver instance (available during execution)."""
+
+    _execution_metadata: "dict[str, Any]" = field(default_factory=dict)
+    """Internal execution metadata for tracking async operations."""
+
+    def __post_init__(self) -> None:
+        """Initialize metadata and extension config if not provided."""
+        if not self.metadata:
+            self.metadata = {}
+        if not self.extension_config:
+            self.extension_config = {}
+
+    @classmethod
+    def from_config(cls, config: Any) -> "MigrationContext":
+        """Create context from database configuration.
+
+        Args:
+            config: Database configuration object.
+
+        Returns:
+            Migration context with dialect information.
+        """
+        dialect = None
+        try:
+            if hasattr(config, "statement_config") and config.statement_config:
+                dialect = getattr(config.statement_config, "dialect", None)
+            elif hasattr(config, "_create_statement_config") and callable(config._create_statement_config):
+                stmt_config = config._create_statement_config()
+                dialect = getattr(stmt_config, "dialect", None)
+        except Exception:
+            logger.debug("Unable to extract dialect from config")
+
+        return cls(dialect=dialect, config=config)
+
+    @property
+    def is_async_execution(self) -> bool:
+        """Check if migrations are running in an async execution context.
+
+        Returns:
+            True if executing in an async context.
+        """
+        try:
+            asyncio.current_task()
+        except RuntimeError:
+            return False
+        else:
+            return True
+
+    @property
+    def is_async_driver(self) -> bool:
+        """Check if the current driver is async.
+
+        Returns:
+            True if driver supports async operations.
+        """
+        if self.driver is None:
+            return False
+
+        execute_method = getattr(self.driver, "execute_script", None)
+        return execute_method is not None and inspect.iscoroutinefunction(execute_method)
+
+    @property
+    def execution_mode(self) -> str:
+        """Get the current execution mode.
+
+        Returns:
+            'async' if in async context, 'sync' otherwise.
+        """
+        return "async" if self.is_async_execution else "sync"
+
+    def set_execution_metadata(self, key: str, value: Any) -> None:
+        """Set execution metadata for tracking migration state.
+
+        Args:
+            key: Metadata key.
+            value: Metadata value.
+        """
+        self._execution_metadata[key] = value
+
+    def get_execution_metadata(self, key: str, default: Any = None) -> Any:
+        """Get execution metadata.
+
+        Args:
+            key: Metadata key.
+            default: Default value if key not found.
+
+        Returns:
+            Metadata value or default.
+        """
+        return self._execution_metadata.get(key, default)
+
+    def validate_async_usage(self, migration_func: Any) -> None:
+        """Validate proper usage of async functions in migration context.
+
+        Args:
+            migration_func: The migration function to validate.
+        """
+        if inspect.iscoroutinefunction(migration_func) and not self.is_async_execution and not self.is_async_driver:
+            msg = (
+                "Async migration function detected but execution context is sync. "
+                "Consider using async database configuration or sync migration functions."
+            )
+            logger.warning(msg)
+
+        if not inspect.iscoroutinefunction(migration_func) and self.is_async_driver:
+            self.set_execution_metadata("mixed_execution", value=True)
+            logger.debug("Sync migration function in async driver context - using compatibility mode")

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,30 +157,49 @@ 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):
     """Loader for Python migration files."""
 
-    __slots__ = ("migrations_dir", "project_root")
+    __slots__ = ("context", "migrations_dir", "project_root")
 
-    def __init__(self, migrations_dir: Path, project_root: "Optional[Path]" = None) -> None:
+    def __init__(self, migrations_dir: Path, project_root: "Path | None" = None, context: "Any | None" = None) -> None:
         """Initialize Python file loader.
 
         Args:
             migrations_dir: Directory containing migration files.
             project_root: Optional project root directory for imports.
+            context: Optional migration context to pass to functions.
         """
         self.migrations_dir = migrations_dir
         self.project_root = project_root if project_root is not None else self._find_project_root(migrations_dir)
+        self.context = context
 
     async def get_up_sql(self, path: Path) -> list[str]:
         """Load Python migration and execute upgrade function.
@@ -208,10 +233,16 @@ class PythonFileLoader(BaseMigrationLoader):
                 msg = f"'{func_name}' is not callable in {path}"
                 raise MigrationLoadError(msg)
 
+            # Check if function accepts context parameter
+            sig = inspect.signature(upgrade_func)
+            accepts_context = "context" in sig.parameters or len(sig.parameters) > 0
+
             if inspect.iscoroutinefunction(upgrade_func):
-                sql_result = await upgrade_func()
+                sql_result = (
+                    await upgrade_func(self.context) if accepts_context and self.context else await upgrade_func()
+                )
             else:
-                sql_result = upgrade_func()
+                sql_result = upgrade_func(self.context) if accepts_context and self.context else upgrade_func()
 
             return self._normalize_and_validate_sql(sql_result, path)
 
@@ -239,10 +270,16 @@ class PythonFileLoader(BaseMigrationLoader):
             if not callable(downgrade_func):
                 return []
 
+            # Check if function accepts context parameter
+            sig = inspect.signature(downgrade_func)
+            accepts_context = "context" in sig.parameters or len(sig.parameters) > 0
+
             if inspect.iscoroutinefunction(downgrade_func):
-                sql_result = await downgrade_func()
+                sql_result = (
+                    await downgrade_func(self.context) if accepts_context and self.context else await downgrade_func()
+                )
             else:
-                sql_result = downgrade_func()
+                sql_result = downgrade_func(self.context) if accepts_context and self.context else downgrade_func()
 
             return self._normalize_and_validate_sql(sql_result, path)
 
@@ -380,7 +417,11 @@ class PythonFileLoader(BaseMigrationLoader):
 
 
 def get_migration_loader(
-    file_path: Path, migrations_dir: Path, project_root: "Optional[Path]" = 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.
 
@@ -388,6 +429,10 @@ def get_migration_loader(
         file_path: Path to the migration file.
         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.
@@ -398,8 +443,8 @@ def get_migration_loader(
     suffix = file_path.suffix
 
     if suffix == ".py":
-        return PythonFileLoader(migrations_dir, project_root)
+        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)