sqlspec 0.17.1__py3-none-any.whl → 0.18.0__py3-none-any.whl

sqlspec/migrations/__init__.py

@@ -1,7 +1,7 @@
 """SQLSpec Migration Tool.
 
 A native migration system for SQLSpec that leverages the SQLFileLoader
-and driver architecture for database versioning.
+and driver system for database versioning.
 """
 
 from sqlspec.migrations.commands import AsyncMigrationCommands, MigrationCommands, SyncMigrationCommands

sqlspec/migrations/base.py

@@ -3,18 +3,19 @@
 This module provides abstract base classes for migration components.
 """
 
+import hashlib
 import operator
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Generic, Optional, TypeVar
+from typing import Any, Generic, Optional, TypeVar, cast
 
 from sqlspec._sql import sql
+from sqlspec.builder import Delete, Insert, Select
 from sqlspec.builder._ddl import CreateTable
-from sqlspec.core.statement import SQL
 from sqlspec.loader import SQLFileLoader
 from sqlspec.migrations.loaders import get_migration_loader
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.sync_tools import run_
+from sqlspec.utils.sync_tools import await_
 
 __all__ = ("BaseMigrationCommands", "BaseMigrationRunner", "BaseMigrationTracker")
 
@@ -28,6 +29,8 @@ ConfigT = TypeVar("ConfigT")
 class BaseMigrationTracker(ABC, Generic[DriverT]):
     """Base class for migration version tracking."""
 
+    __slots__ = ("version_table",)
+
     def __init__(self, version_table_name: str = "ddl_migrations") -> None:
         """Initialize the migration tracker.
 
@@ -36,54 +39,43 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
         """
         self.version_table = version_table_name
 
-    def _get_create_table_sql(self) -> SQL:
-        """Get SQL for creating the tracking table.
+    def _get_create_table_sql(self) -> CreateTable:
+        """Get SQL builder for creating the tracking table.
 
         Returns:
-            SQL object for table creation.
+            SQL builder object for table creation.
         """
-        builder = CreateTable(self.version_table)
-        if not hasattr(builder, "_columns"):
-            builder._columns = []
-        if not hasattr(builder, "_constraints"):
-            builder._constraints = []
-        if not hasattr(builder, "_table_options"):
-            builder._table_options = {}
-
         return (
-            builder.if_not_exists()
+            sql.create_table(self.version_table)
+            .if_not_exists()
             .column("version_num", "VARCHAR(32)", primary_key=True)
             .column("description", "TEXT")
-            .column("applied_at", "TIMESTAMP", not_null=True, default="CURRENT_TIMESTAMP")
+            .column("applied_at", "TIMESTAMP", default="CURRENT_TIMESTAMP", not_null=True)
             .column("execution_time_ms", "INTEGER")
             .column("checksum", "VARCHAR(64)")
             .column("applied_by", "VARCHAR(255)")
-        ).to_statement()
+        )
 
-    def _get_current_version_sql(self) -> SQL:
-        """Get SQL for retrieving current version.
+    def _get_current_version_sql(self) -> Select:
+        """Get SQL builder for retrieving current version.
 
         Returns:
-            SQL object for version query.
+            SQL builder object for version query.
         """
+        return sql.select("version_num").from_(self.version_table).order_by("version_num DESC").limit(1)
 
-        return (
-            sql.select("version_num").from_(self.version_table).order_by("version_num DESC").limit(1)
-        ).to_statement()
-
-    def _get_applied_migrations_sql(self) -> SQL:
-        """Get SQL for retrieving all applied migrations.
+    def _get_applied_migrations_sql(self) -> Select:
+        """Get SQL builder for retrieving all applied migrations.
 
         Returns:
-            SQL object for migrations query.
+            SQL builder object for migrations query.
         """
-
-        return (sql.select("*").from_(self.version_table).order_by("version_num")).to_statement()
+        return sql.select("*").from_(self.version_table).order_by("version_num")
 
     def _get_record_migration_sql(
         self, version: str, description: str, execution_time_ms: int, checksum: str, applied_by: str
-    ) -> SQL:
-        """Get SQL for recording a migration.
+    ) -> Insert:
+        """Get SQL builder for recording a migration.
 
         Args:
             version: Version number of the migration.
@@ -93,26 +85,24 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
             applied_by: User who applied the migration.
 
         Returns:
-            SQL object for insert.
+            SQL builder object for insert.
         """
-
         return (
             sql.insert(self.version_table)
             .columns("version_num", "description", "execution_time_ms", "checksum", "applied_by")
             .values(version, description, execution_time_ms, checksum, applied_by)
-        ).to_statement()
+        )
 
-    def _get_remove_migration_sql(self, version: str) -> SQL:
-        """Get SQL for removing a migration record.
+    def _get_remove_migration_sql(self, version: str) -> Delete:
+        """Get SQL builder for removing a migration record.
 
         Args:
             version: Version number to remove.
 
         Returns:
-            SQL object for delete.
+            SQL builder object for delete.
         """
-
-        return (sql.delete().from_(self.version_table).where(sql.version_num == version)).to_statement()
+        return sql.delete().from_(self.version_table).where(sql.version_num == version)
 
     @abstractmethod
     def ensure_tracking_table(self, driver: DriverT) -> Any:
@@ -176,7 +166,6 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         Returns:
             MD5 checksum hex string.
         """
-        import hashlib
 
         return hashlib.md5(content.encode()).hexdigest()  # noqa: S324
 
@@ -226,7 +215,7 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
             has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
         else:
             try:
-                has_downgrade = bool(run_(loader.get_down_sql)(file_path))
+                has_downgrade = bool(await_(loader.get_down_sql, raise_sync_error=False)(file_path))
             except Exception:
                 has_downgrade = False
 
@@ -240,7 +229,7 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
             "loader": loader,
         }
 
-    def _get_migration_sql(self, migration: "dict[str, Any]", direction: str) -> Optional[SQL]:
+    def _get_migration_sql(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
         """Get migration SQL for given direction.
 
         Args:
@@ -261,7 +250,7 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
 
         try:
             method = loader.get_up_sql if direction == "up" else loader.get_down_sql
-            sql_statements = run_(method)(file_path)
+            sql_statements = await_(method, raise_sync_error=False)(file_path)
 
         except Exception as e:
             if direction == "down":
@@ -271,7 +260,7 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
             raise ValueError(msg) from e
         else:
             if sql_statements:
-                return SQL(sql_statements[0])
+                return cast("list[str]", sql_statements)
             return None
 
     @abstractmethod
@@ -312,7 +301,7 @@ class BaseMigrationCommands(ABC, Generic[ConfigT, DriverT]):
         self.config = config
         migration_config = getattr(self.config, "migration_config", {}) or {}
 
-        self.version_table = migration_config.get("version_table_name", "sqlspec_migrations")
+        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
 

sqlspec/migrations/commands.py

@@ -3,7 +3,7 @@
 This module provides the main command interface for database migrations.
 """
 
-from typing import TYPE_CHECKING, Any, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 from rich.console import Console
 from rich.table import Table
@@ -11,7 +11,6 @@ from rich.table import Table
 from sqlspec._sql import sql
 from sqlspec.migrations.base import BaseMigrationCommands
 from sqlspec.migrations.runner import AsyncMigrationRunner, SyncMigrationRunner
-from sqlspec.migrations.tracker import AsyncMigrationTracker, SyncMigrationTracker
 from sqlspec.migrations.utils import create_migration_file
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.sync_tools import await_
@@ -26,7 +25,7 @@ console = Console()
 
 
 class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
-    """SQLSpec native migration commands."""
+    """Synchronous migration commands."""
 
     def __init__(self, config: "SyncConfigT") -> None:
         """Initialize migration commands.
@@ -35,7 +34,7 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
             config: The SQLSpec configuration.
         """
         super().__init__(config)
-        self.tracker = SyncMigrationTracker(self.version_table)
+        self.tracker = config.migration_tracker_type(self.version_table)
         self.runner = SyncMigrationRunner(self.migrations_path)
 
     def init(self, directory: str, package: bool = True) -> None:
@@ -47,11 +46,14 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
         """
         self.init_directory(directory, package)
 
-    def current(self, verbose: bool = False) -> None:
+    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.
         """
         with self.config.provide_session() as driver:
             self.tracker.ensure_tracking_table(driver)
@@ -59,7 +61,7 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
             current = self.tracker.get_current_version(driver)
             if not current:
                 console.print("[yellow]No migrations applied yet[/]")
-                return
+                return None
 
             console.print(f"[green]Current version:[/] {current}")
 
@@ -84,6 +86,8 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
 
                 console.print(table)
 
+            return cast("Optional[str]", current)
+
     def upgrade(self, revision: str = "head") -> None:
         """Upgrade to a target revision.
 
@@ -137,6 +141,8 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
             to_revert = []
             if revision == "-1":
                 to_revert = [applied[-1]]
+            elif revision == "base":
+                to_revert = list(reversed(applied))
             else:
                 for migration in reversed(applied):
                     if migration["version_num"] > revision:
@@ -195,7 +201,7 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
 
 
 class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
-    """SQLSpec native migration commands."""
+    """Asynchronous migration commands."""
 
     def __init__(self, sqlspec_config: "AsyncConfigT") -> None:
         """Initialize migration commands.
@@ -204,7 +210,7 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
             sqlspec_config: The SQLSpec configuration.
         """
         super().__init__(sqlspec_config)
-        self.tracker = AsyncMigrationTracker(self.version_table)
+        self.tracker = sqlspec_config.migration_tracker_type(self.version_table)
         self.runner = AsyncMigrationRunner(self.migrations_path)
 
     async def init(self, directory: str, package: bool = True) -> None:
@@ -216,11 +222,14 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
         """
         self.init_directory(directory, package)
 
-    async def current(self, verbose: bool = False) -> None:
+    async 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.
         """
         async with self.config.provide_session() as driver:
             await self.tracker.ensure_tracking_table(driver)
@@ -228,7 +237,7 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
             current = await self.tracker.get_current_version(driver)
             if not current:
                 console.print("[yellow]No migrations applied yet[/]")
-                return
+                return None
 
             console.print(f"[green]Current version:[/] {current}")
             if verbose:
@@ -249,6 +258,8 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
                     )
                 console.print(table)
 
+            return cast("Optional[str]", current)
+
     async def upgrade(self, revision: str = "head") -> None:
         """Upgrade to a target revision.
 
@@ -297,6 +308,8 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
             to_revert = []
             if revision == "-1":
                 to_revert = [applied[-1]]
+            elif revision == "base":
+                to_revert = list(reversed(applied))
             else:
                 for migration in reversed(applied):
                     if migration["version_num"] > revision:
@@ -382,20 +395,26 @@ class MigrationCommands:
             package: Whether to create __init__.py file.
         """
         if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).init)(directory, package=package)
+            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) -> None:
+    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:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).current, raise_sync_error=False)(verbose=verbose)
-        else:
-            cast("SyncMigrationCommands[Any]", self._impl).current(verbose=verbose)
+            return await_(cast("AsyncMigrationCommands[Any]", self._impl).current, raise_sync_error=False)(
+                verbose=verbose
+            )
+        return cast("SyncMigrationCommands[Any]", self._impl).current(verbose=verbose)
 
     def upgrade(self, revision: str = "head") -> None:
         """Upgrade to a target revision.

sqlspec/migrations/loaders.py

@@ -70,7 +70,7 @@ class BaseMigrationLoader(abc.ABC):
 
 
 class SQLFileLoader(BaseMigrationLoader):
-    """Loader for SQL migration files using SQLFileLoader."""
+    """Loader for SQL migration files."""
 
     __slots__ = ("sql_loader",)
 

sqlspec/migrations/runner.py

@@ -5,13 +5,13 @@ This module handles migration file loading and execution using SQLFileLoader.
 
 import time
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from sqlspec.core.statement import SQL
 from sqlspec.migrations.base import BaseMigrationRunner
 from sqlspec.migrations.loaders import get_migration_loader
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.sync_tools import run_
+from sqlspec.utils.sync_tools import await_
 
 if TYPE_CHECKING:
     from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
@@ -22,7 +22,7 @@ logger = get_logger("migrations.runner")
 
 
 class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
-    """Executes migrations using SQLFileLoader."""
+    """Synchronous migration executor."""
 
     def get_migration_files(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
@@ -55,12 +55,15 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        upgrade_sql = self._get_migration_sql(migration, "up")
-        if upgrade_sql is None:
+        upgrade_sql_list = self._get_migration_sql(migration, "up")
+        if upgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        driver.execute(upgrade_sql)
+
+        for sql_statement in upgrade_sql_list:
+            if sql_statement.strip():
+                driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -76,12 +79,15 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        downgrade_sql = self._get_migration_sql(migration, "down")
-        if downgrade_sql is None:
+        downgrade_sql_list = self._get_migration_sql(migration, "down")
+        if downgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        driver.execute(downgrade_sql)
+
+        for sql_statement in downgrade_sql_list:
+            if sql_statement.strip():
+                driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -103,8 +109,8 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
                 loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
 
                 try:
-                    up_sql = run_(loader.get_up_sql)(file_path)
-                    down_sql = run_(loader.get_down_sql)(file_path)
+                    up_sql = await_(loader.get_up_sql, raise_sync_error=False)(file_path)
+                    down_sql = await_(loader.get_down_sql, raise_sync_error=False)(file_path)
 
                     if up_sql:
                         all_queries[f"migrate-{version}-up"] = SQL(up_sql[0])
@@ -118,7 +124,7 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
 
 
 class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
-    """Executes migrations using SQLFileLoader."""
+    """Asynchronous migration executor."""
 
     async def get_migration_files(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
@@ -137,7 +143,80 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Dictionary containing migration metadata.
         """
-        return self._load_migration_metadata(file_path)
+        return await self._load_migration_metadata_async(file_path)
+
+    async def _load_migration_metadata_async(self, file_path: Path) -> "dict[str, Any]":
+        """Load migration metadata from file (async version).
+
+        Args:
+            file_path: Path to the migration file.
+
+        Returns:
+            Migration metadata dictionary.
+        """
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
+        loader.validate_migration_file(file_path)
+        content = file_path.read_text(encoding="utf-8")
+        checksum = self._calculate_checksum(content)
+        version = self._extract_version(file_path.name)
+        description = file_path.stem.split("_", 1)[1] if "_" in file_path.stem else ""
+
+        has_upgrade, has_downgrade = True, False
+
+        if file_path.suffix == ".sql":
+            up_query, down_query = f"migrate-{version}-up", f"migrate-{version}-down"
+            self.loader.clear_cache()
+            self.loader.load_sql(file_path)
+            has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
+        else:
+            try:
+                has_downgrade = bool(await loader.get_down_sql(file_path))
+            except Exception:
+                has_downgrade = False
+
+        return {
+            "version": version,
+            "description": description,
+            "file_path": file_path,
+            "checksum": checksum,
+            "has_upgrade": has_upgrade,
+            "has_downgrade": has_downgrade,
+            "loader": loader,
+        }
+
+    async def _get_migration_sql_async(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
+        """Get migration SQL for given direction (async version).
+
+        Args:
+            migration: Migration metadata.
+            direction: Either 'up' or 'down'.
+
+        Returns:
+            SQL statements for the migration.
+        """
+        if not migration.get(f"has_{direction}grade"):
+            if direction == "down":
+                logger.warning("Migration %s has no downgrade query", migration["version"])
+                return None
+            msg = f"Migration {migration['version']} has no upgrade query"
+            raise ValueError(msg)
+
+        file_path, loader = migration["file_path"], migration["loader"]
+
+        try:
+            method = loader.get_up_sql if direction == "up" else loader.get_down_sql
+            sql_statements = await method(file_path)
+
+        except Exception as e:
+            if direction == "down":
+                logger.warning("Failed to load downgrade for migration %s: %s", migration["version"], e)
+                return None
+            msg = f"Failed to load upgrade for migration {migration['version']}: {e}"
+            raise ValueError(msg) from e
+        else:
+            if sql_statements:
+                return cast("list[str]", sql_statements)
+            return None
 
     async def execute_upgrade(
         self, driver: "AsyncDriverAdapterBase", migration: "dict[str, Any]"
@@ -151,12 +230,15 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        upgrade_sql = self._get_migration_sql(migration, "up")
-        if upgrade_sql is None:
+        upgrade_sql_list = await self._get_migration_sql_async(migration, "up")
+        if upgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        await driver.execute(upgrade_sql)
+
+        for sql_statement in upgrade_sql_list:
+            if sql_statement.strip():
+                await driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -172,12 +254,15 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        downgrade_sql = self._get_migration_sql(migration, "down")
-        if downgrade_sql is None:
+        downgrade_sql_list = await self._get_migration_sql_async(migration, "down")
+        if downgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        await driver.execute(downgrade_sql)
+
+        for sql_statement in downgrade_sql_list:
+            if sql_statement.strip():
+                await driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time