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

sqlspec/migrations/base.py

@@ -15,6 +15,7 @@ from sqlspec.builder._ddl import CreateTable
 from sqlspec.loader import SQLFileLoader
 from sqlspec.migrations.loaders import get_migration_loader
 from sqlspec.utils.logging import get_logger
+from sqlspec.utils.module_loader import module_to_os_path
 from sqlspec.utils.sync_tools import await_
 
 __all__ = ("BaseMigrationCommands", "BaseMigrationRunner", "BaseMigrationTracker")
@@ -135,15 +136,29 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
 class BaseMigrationRunner(ABC, Generic[DriverT]):
     """Base class for migration execution."""
 
-    def __init__(self, migrations_path: Path) -> None:
+    extension_configs: "dict[str, dict[str, Any]]"
+
+    def __init__(
+        self,
+        migrations_path: Path,
+        extension_migrations: "Optional[dict[str, Path]]" = None,
+        context: "Optional[Any]" = None,
+        extension_configs: "Optional[dict[str, dict[str, Any]]]" = None,
+    ) -> 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.
         """
         self.migrations_path = migrations_path
+        self.extension_migrations = extension_migrations or {}
         self.loader = SQLFileLoader()
         self.project_root: Optional[Path] = None
+        self.context = context
+        self.extension_configs = extension_configs or {}
 
     def _extract_version(self, filename: str) -> Optional[str]:
         """Extract version from filename.
@@ -154,6 +169,12 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         Returns:
             The extracted version string or None.
         """
+        # Handle extension-prefixed versions (e.g., "ext_litestar_0001")
+        if filename.startswith("ext_"):
+            # This is already a prefixed version, return as-is
+            return filename
+
+        # Regular version extraction
         parts = filename.split("_", 1)
         return parts[0].zfill(4) if parts and parts[0].isdigit() else None
 
@@ -175,17 +196,31 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         Returns:
             List of tuples containing (version, file_path).
         """
-        if not self.migrations_path.exists():
-            return []
-
         migrations = []
-        for pattern in ["*.sql", "*.py"]:
-            for file_path in self.migrations_path.glob(pattern):
-                if file_path.name.startswith("."):
-                    continue
-                version = self._extract_version(file_path.name)
-                if version:
-                    migrations.append((version, file_path))
+
+        # Scan primary migration path
+        if self.migrations_path.exists():
+            for pattern in ("*.sql", "*.py"):
+                for file_path in self.migrations_path.glob(pattern):
+                    if file_path.name.startswith("."):
+                        continue
+                    version = self._extract_version(file_path.name)
+                    if version:
+                        migrations.append((version, file_path))
+
+        # Scan extension migration paths
+        for ext_name, ext_path in self.extension_migrations.items():
+            if ext_path.exists():
+                for pattern in ("*.sql", "*.py"):
+                    for file_path in ext_path.glob(pattern):
+                        if file_path.name.startswith("."):
+                            continue
+                        # Prefix extension migrations to avoid version conflicts
+                        version = self._extract_version(file_path.name)
+                        if version:
+                            # Use ext_ prefix to distinguish extension migrations
+                            prefixed_version = f"ext_{ext_name}_{version}"
+                            migrations.append((prefixed_version, file_path))
 
         return sorted(migrations, key=operator.itemgetter(0))
 
@@ -199,7 +234,45 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
             Migration metadata dictionary.
         """
 
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
+        # Check if this is an extension migration and update context accordingly
+        context_to_use = self.context
+        if context_to_use and file_path.name.startswith("ext_"):
+            # Try to extract extension name from the version
+            version = self._extract_version(file_path.name)
+            if version and version.startswith("ext_"):
+                # Parse extension name from version like "ext_litestar_0001"
+                min_extension_version_parts = 3
+                parts = version.split("_", 2)
+                if len(parts) >= min_extension_version_parts:
+                    ext_name = parts[1]
+                    if ext_name in self.extension_configs:
+                        # Create a new context with the extension config
+                        from sqlspec.migrations.context import MigrationContext
+
+                        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 extension migrations, check by path
+        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:
+                    from sqlspec.migrations.context import MigrationContext
+
+                    context_to_use = MigrationContext(
+                        dialect=self.context.dialect,
+                        config=self.context.config,
+                        driver=self.context.driver,
+                        metadata=self.context.metadata.copy() if self.context.metadata else {},
+                        extension_config=self.extension_configs[ext_name],
+                    )
+                break
+
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use)
         loader.validate_migration_file(file_path)
         content = file_path.read_text(encoding="utf-8")
         checksum = self._calculate_checksum(content)
@@ -292,6 +365,8 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
 class BaseMigrationCommands(ABC, Generic[ConfigT, DriverT]):
     """Base class for migration commands."""
 
+    extension_configs: "dict[str, dict[str, Any]]"
+
     def __init__(self, config: ConfigT) -> None:
         """Initialize migration commands.
 
@@ -304,6 +379,72 @@ class BaseMigrationCommands(ABC, Generic[ConfigT, DriverT]):
         self.version_table = migration_config.get("version_table_name", "ddl_migrations")
         self.migrations_path = Path(migration_config.get("script_location", "migrations"))
         self.project_root = Path(migration_config["project_root"]) if "project_root" in migration_config else None
+        self.include_extensions = migration_config.get("include_extensions", [])
+        self.extension_configs = self._parse_extension_configs()
+
+    def _parse_extension_configs(self) -> "dict[str, dict[str, Any]]":
+        """Parse extension configurations from include_extensions.
+
+        Supports both string format (extension name) and dict format
+        (extension name with configuration).
+
+        Returns:
+            Dictionary mapping extension names to their configurations.
+        """
+        configs = {}
+
+        for ext_config in self.include_extensions:
+            if isinstance(ext_config, str):
+                # Simple string format: just the extension name
+                ext_name = ext_config
+                ext_options = {}
+            elif isinstance(ext_config, dict):
+                # Dict format: {"name": "litestar", "session_table": "custom_sessions"}
+                ext_name_raw = ext_config.get("name")
+                if not ext_name_raw:
+                    logger.warning("Extension configuration missing 'name' field: %s", ext_config)
+                    continue
+                # Assert for type narrowing: ext_name_raw is guaranteed to be str here
+                assert isinstance(ext_name_raw, str)
+                ext_name = ext_name_raw
+                ext_options = {k: v for k, v in ext_config.items() if k != "name"}
+            else:
+                logger.warning("Invalid extension configuration format: %s", ext_config)
+                continue
+
+            # Apply default configurations for known extensions
+            if ext_name == "litestar" and "session_table" not in ext_options:
+                ext_options["session_table"] = "litestar_sessions"
+
+            configs[ext_name] = ext_options
+
+        return configs
+
+    def _discover_extension_migrations(self) -> "dict[str, Path]":
+        """Discover migration paths for configured extensions.
+
+        Returns:
+            Dictionary mapping extension names to their migration paths.
+        """
+
+        extension_migrations = {}
+
+        for ext_name in self.extension_configs:
+            module_name = "sqlspec.extensions.litestar" if ext_name == "litestar" else f"sqlspec.extensions.{ext_name}"
+
+            try:
+                module_path = module_to_os_path(module_name)
+                migrations_dir = module_path / "migrations"
+
+                if migrations_dir.exists():
+                    extension_migrations[ext_name] = migrations_dir
+                    logger.debug("Found migrations for extension %s at %s", ext_name, migrations_dir)
+                else:
+                    logger.warning("No migrations directory found for extension %s", ext_name)
+            except TypeError:
+                logger.warning("Extension %s not found", ext_name)
+
+        return extension_migrations
 
     def _get_init_readme_content(self) -> str:
         """Get README content for migration directory initialization.
@@ -368,8 +509,6 @@ This naming ensures proper sorting and avoids conflicts when loading multiple fi
         readme = migrations_dir / "README.md"
         readme.write_text(self._get_init_readme_content())
 
-        (migrations_dir / ".gitkeep").touch()
-
         console.print(f"[green]Initialized migrations in {directory}[/]")
 
     @abstractmethod

sqlspec/migrations/commands.py

@@ -10,15 +10,15 @@ from rich.table import Table
 
 from sqlspec._sql import sql
 from sqlspec.migrations.base import BaseMigrationCommands
+from sqlspec.migrations.context import MigrationContext
 from sqlspec.migrations.runner import AsyncMigrationRunner, SyncMigrationRunner
 from sqlspec.migrations.utils import create_migration_file
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.sync_tools import await_
 
 if TYPE_CHECKING:
     from sqlspec.config import AsyncConfigT, SyncConfigT
 
-__all__ = ("AsyncMigrationCommands", "MigrationCommands", "SyncMigrationCommands")
+__all__ = ("AsyncMigrationCommands", "SyncMigrationCommands", "create_migration_commands")
 
 logger = get_logger("migrations.commands")
 console = Console()
@@ -35,7 +35,14 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
         """
         super().__init__(config)
         self.tracker = config.migration_tracker_type(self.version_table)
-        self.runner = SyncMigrationRunner(self.migrations_path)
+
+        # Create context with extension configurations
+        context = MigrationContext.from_config(config)
+        context.extension_config = self.extension_configs
+
+        self.runner = SyncMigrationRunner(
+            self.migrations_path, self._discover_extension_migrations(), context, self.extension_configs
+        )
 
     def init(self, directory: str, package: bool = True) -> None:
         """Initialize migration directory structure.
@@ -203,15 +210,22 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
 class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
     """Asynchronous migration commands."""
 
-    def __init__(self, sqlspec_config: "AsyncConfigT") -> None:
+    def __init__(self, config: "AsyncConfigT") -> None:
         """Initialize migration commands.
 
         Args:
-            sqlspec_config: The SQLSpec configuration.
+            config: The SQLSpec configuration.
         """
-        super().__init__(sqlspec_config)
-        self.tracker = sqlspec_config.migration_tracker_type(self.version_table)
-        self.runner = AsyncMigrationRunner(self.migrations_path)
+        super().__init__(config)
+        self.tracker = config.migration_tracker_type(self.version_table)
+
+        # Create context with extension configurations
+        context = MigrationContext.from_config(config)
+        context.extension_config = self.extension_configs
+
+        self.runner = AsyncMigrationRunner(
+            self.migrations_path, self._discover_extension_migrations(), context, self.extension_configs
+        )
 
     async def init(self, directory: str, package: bool = True) -> None:
         """Initialize migration directory structure.
@@ -370,93 +384,17 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
         console.print(f"[green]Created migration:[/] {file_path}")
 
 
-class MigrationCommands:
-    """Unified migration commands that adapt to sync/async configs."""
-
-    def __init__(self, config: "Union[SyncConfigT, AsyncConfigT]") -> None:
-        """Initialize migration commands with sync/async implementation.
-
-        Args:
-            config: The SQLSpec configuration.
-        """
-        if config.is_async:
-            self._impl: Union[AsyncMigrationCommands[Any], SyncMigrationCommands[Any]] = AsyncMigrationCommands(
-                cast("AsyncConfigT", config)
-            )
-        else:
-            self._impl = SyncMigrationCommands(cast("SyncConfigT", config))
-        self._is_async = config.is_async
-
-    def init(self, directory: str, package: bool = True) -> None:
-        """Initialize migration directory structure.
+def create_migration_commands(
+    config: "Union[SyncConfigT, AsyncConfigT]",
+) -> "Union[SyncMigrationCommands[Any], AsyncMigrationCommands[Any]]":
+    """Factory function to create the appropriate migration commands.
 
-        Args:
-            directory: Directory to initialize migrations in.
-            package: Whether to create __init__.py file.
-        """
-        if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).init, raise_sync_error=False)(
-                directory, package=package
-            )
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).init(directory, package=package)
-
-    def current(self, verbose: bool = False) -> "Optional[str]":
-        """Show current migration version.
-
-        Args:
-            verbose: Whether to show detailed migration history.
-
-        Returns:
-            The current migration version or None if no migrations applied.
-        """
-        if self._is_async:
-            return await_(cast("AsyncMigrationCommands[Any]", self._impl).current, raise_sync_error=False)(
-                verbose=verbose
-            )
-        return cast("SyncMigrationCommands[Any]", self._impl).current(verbose=verbose)
+    Args:
+        config: The SQLSpec configuration.
 
-    def upgrade(self, revision: str = "head") -> None:
-        """Upgrade to a target revision.
-
-        Args:
-            revision: Target revision or "head" for latest.
-        """
-        if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).upgrade, raise_sync_error=False)(revision=revision)
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).upgrade(revision=revision)
-
-    def downgrade(self, revision: str = "-1") -> None:
-        """Downgrade to a target revision.
-
-        Args:
-            revision: Target revision or "-1" for one step back.
-        """
-        if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).downgrade, raise_sync_error=False)(revision=revision)
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).downgrade(revision=revision)
-
-    def stamp(self, revision: str) -> None:
-        """Mark database as being at a specific revision without running migrations.
-
-        Args:
-            revision: The revision to stamp.
-        """
-        if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).stamp, raise_sync_error=False)(revision)
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).stamp(revision)
-
-    def revision(self, message: str, file_type: str = "sql") -> None:
-        """Create a new migration file.
-
-        Args:
-            message: Description for the migration.
-            file_type: Type of migration file to create ('sql' or 'py').
-        """
-        if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).revision, raise_sync_error=False)(message, file_type)
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).revision(message, file_type)
+    Returns:
+        Appropriate migration commands instance.
+    """
+    if config.is_async:
+        return AsyncMigrationCommands(cast("AsyncConfigT", config))
+    return SyncMigrationCommands(cast("SyncConfigT", config))

sqlspec/migrations/context.py

@@ -0,0 +1,145 @@
+"""Migration context for passing runtime information to migrations."""
+
+import asyncio
+import inspect
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+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: "Optional[Any]" = None
+    """Database configuration object."""
+    dialect: "Optional[str]" = None
+    """Database dialect (e.g., 'postgres', 'mysql', 'sqlite')."""
+    metadata: "Optional[dict[str, Any]]" = None
+    """Additional metadata for the migration."""
+    extension_config: "Optional[dict[str, Any]]" = None
+    """Extension-specific configuration options."""
+
+    driver: "Optional[Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]]" = 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.
+
+        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 = (
+                "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/loaders.py

@@ -164,17 +164,21 @@ class SQLFileLoader(BaseMigrationLoader):
 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: "Optional[Path]" = None, context: "Optional[Any]" = 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 +212,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 +249,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 +396,7 @@ 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: "Optional[Path]" = None, context: "Optional[Any]" = None
 ) -> BaseMigrationLoader:
     """Factory function to get appropriate loader for migration file.
 
@@ -388,6 +404,7 @@ 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.
 
     Returns:
         Appropriate loader instance for the file type.
@@ -398,7 +415,7 @@ 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()
     msg = f"Unsupported migration file type: {suffix}"