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

sqlspec/migrations/runner.py

@@ -4,11 +4,11 @@ This module provides separate sync and async migration runners with clean separa
 of concerns and proper type safety.
 """
 
-import operator
+import inspect
 import time
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Literal, Union, cast, overload
 
 from sqlspec.core.statement import SQL
 from sqlspec.migrations.context import MigrationContext
@@ -17,7 +17,7 @@ from sqlspec.utils.logging import get_logger
 from sqlspec.utils.sync_tools import async_, await_
 
 if TYPE_CHECKING:
-    from collections.abc import Coroutine
+    from collections.abc import Awaitable, Callable, Coroutine
 
     from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
 
@@ -32,9 +32,9 @@ class BaseMigrationRunner(ABC):
     def __init__(
         self,
         migrations_path: Path,
-        extension_migrations: "Optional[dict[str, Path]]" = None,
-        context: "Optional[MigrationContext]" = None,
-        extension_configs: "Optional[dict[str, dict[str, Any]]]" = None,
+        extension_migrations: "dict[str, Path] | None" = None,
+        context: "MigrationContext | None" = None,
+        extension_configs: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize the migration runner.
 
@@ -49,31 +49,46 @@ class BaseMigrationRunner(ABC):
         from sqlspec.loader import SQLFileLoader
 
         self.loader = SQLFileLoader()
-        self.project_root: Optional[Path] = None
+        self.project_root: Path | None = None
         self.context = context
         self.extension_configs = extension_configs or {}
 
-    def _extract_version(self, filename: str) -> "Optional[str]":
+    def _extract_version(self, filename: str) -> "str | None":
         """Extract version from filename.
 
+        Supports sequential (0001), timestamp (20251011120000), and extension-prefixed
+        (ext_litestar_0001) version formats.
+
         Args:
             filename: The migration filename.
 
         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
+        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 None
 
-        # Regular version extraction
-        parts = filename.split("_", 1)
-        return parts[0].zfill(4) if parts and parts[0].isdigit() else None
+        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 None
 
     def _calculate_checksum(self, content: str) -> str:
         """Calculate MD5 checksum of migration content.
 
+        Canonicalizes content by excluding query name headers that change during
+        fix command (migrate-{version}-up/down). This ensures checksums remain
+        stable when converting timestamp versions to sequential format.
+
         Args:
             content: The migration file content.
 
@@ -81,8 +96,11 @@ class BaseMigrationRunner(ABC):
             MD5 checksum hex string.
         """
         import hashlib
+        import re
+
+        canonical_content = re.sub(r"^--\s*name:\s*migrate-[^-]+-(?:up|down)\s*$", "", content, flags=re.MULTILINE)
 
-        return hashlib.md5(content.encode()).hexdigest()  # noqa: S324
+        return hashlib.md5(canonical_content.encode()).hexdigest()  # noqa: S324
 
     @abstractmethod
     def load_migration(self, file_path: Path) -> Union["dict[str, Any]", "Coroutine[Any, Any, dict[str, Any]]"]:
@@ -129,7 +147,16 @@ class BaseMigrationRunner(ABC):
                             prefixed_version = f"ext_{ext_name}_{version}"
                             migrations.append((prefixed_version, file_path))
 
-        return sorted(migrations, key=operator.itemgetter(0))
+        from sqlspec.utils.version import parse_version
+
+        def version_sort_key(migration_tuple: "tuple[str, Path]") -> "Any":
+            version_str = migration_tuple[0]
+            try:
+                return parse_version(version_str)
+            except ValueError:
+                return version_str
+
+        return sorted(migrations, key=version_sort_key)
 
     def get_migration_files(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
@@ -139,29 +166,41 @@ class BaseMigrationRunner(ABC):
         """
         return self._get_migration_files_sync()
 
-    def _load_migration_metadata_common(self, file_path: Path) -> "dict[str, Any]":
+    def _load_migration_metadata_common(self, file_path: Path, version: "str | None" = None) -> "dict[str, Any]":
         """Load common migration metadata that doesn't require async operations.
 
         Args:
             file_path: Path to the migration file.
+            version: Optional pre-extracted version (preserves prefixes like ext_adk_0001).
 
         Returns:
             Partial migration metadata dictionary.
         """
+        import re
+
         content = file_path.read_text(encoding="utf-8")
         checksum = self._calculate_checksum(content)
-        version = self._extract_version(file_path.name)
+        if version is None:
+            version = self._extract_version(file_path.name)
         description = file_path.stem.split("_", 1)[1] if "_" in file_path.stem else ""
 
+        transactional_match = re.search(
+            r"^--\s*transactional:\s*(true|false)\s*$", content, re.MULTILINE | re.IGNORECASE
+        )
+        transactional = None
+        if transactional_match:
+            transactional = transactional_match.group(1).lower() == "true"
+
         return {
             "version": version,
             "description": description,
             "file_path": file_path,
             "checksum": checksum,
             "content": content,
+            "transactional": transactional,
         }
 
-    def _get_context_for_migration(self, file_path: Path) -> "Optional[MigrationContext]":
+    def _get_context_for_migration(self, file_path: Path) -> "MigrationContext | None":
         """Get the appropriate context for a migration file.
 
         Args:
@@ -201,24 +240,43 @@ class BaseMigrationRunner(ABC):
 
         return context_to_use
 
+    def should_use_transaction(self, migration: "dict[str, Any]", config: Any) -> bool:
+        """Determine if migration should run in a transaction.
+
+        Args:
+            migration: Migration metadata dictionary.
+            config: The database configuration instance.
+
+        Returns:
+            True if migration should be wrapped in a transaction.
+        """
+        if not config.supports_transactional_ddl:
+            return False
+
+        if migration.get("transactional") is not None:
+            return bool(migration["transactional"])
+
+        migration_config = getattr(config, "migration_config", {}) or {}
+        return bool(migration_config.get("transactional", True))
+
 
 class SyncMigrationRunner(BaseMigrationRunner):
     """Synchronous migration runner with pure sync methods."""
 
-    def load_migration(self, file_path: Path) -> "dict[str, Any]":
+    def load_migration(self, file_path: Path, version: "str | None" = None) -> "dict[str, Any]":
         """Load a migration file and extract its components.
 
         Args:
             file_path: Path to the migration file.
+            version: Optional pre-extracted version (preserves prefixes like ext_adk_0001).
 
         Returns:
             Dictionary containing migration metadata and queries.
         """
-        # Get common metadata
-        metadata = self._load_migration_metadata_common(file_path)
+        metadata = self._load_migration_metadata_common(file_path, version)
         context_to_use = self._get_context_for_migration(file_path)
 
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use)
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, self.loader)
         loader.validate_migration_file(file_path)
 
         has_upgrade, has_downgrade = True, False
@@ -226,8 +284,6 @@ class SyncMigrationRunner(BaseMigrationRunner):
         if file_path.suffix == ".sql":
             version = metadata["version"]
             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:
@@ -239,13 +295,20 @@ class SyncMigrationRunner(BaseMigrationRunner):
         return metadata
 
     def execute_upgrade(
-        self, driver: "SyncDriverAdapterBase", migration: "dict[str, Any]"
-    ) -> "tuple[Optional[str], int]":
+        self,
+        driver: "SyncDriverAdapterBase",
+        migration: "dict[str, Any]",
+        *,
+        use_transaction: "bool | None" = None,
+        on_success: "Callable[[int], None] | None" = None,
+    ) -> "tuple[str | None, int]":
         """Execute an upgrade migration.
 
         Args:
             driver: The sync database driver to use.
             migration: Migration metadata dictionary.
+            use_transaction: Override transaction behavior. If None, uses should_use_transaction logic.
+            on_success: Callback invoked with execution_time_ms before commit (for version tracking).
 
         Returns:
             Tuple of (sql_content, execution_time_ms).
@@ -254,22 +317,50 @@ class SyncMigrationRunner(BaseMigrationRunner):
         if upgrade_sql_list is None:
             return None, 0
 
+        if use_transaction is None:
+            config = self.context.config if self.context else None
+            use_transaction = self.should_use_transaction(migration, config) if config else False
+
         start_time = time.time()
 
-        for sql_statement in upgrade_sql_list:
-            if sql_statement.strip():
-                driver.execute_script(sql_statement)
-        execution_time = int((time.time() - start_time) * 1000)
+        if use_transaction:
+            try:
+                driver.begin()
+                for sql_statement in upgrade_sql_list:
+                    if sql_statement.strip():
+                        driver.execute_script(sql_statement)
+                execution_time = int((time.time() - start_time) * 1000)
+                if on_success:
+                    on_success(execution_time)
+                driver.commit()
+            except Exception:
+                driver.rollback()
+                raise
+        else:
+            for sql_statement in upgrade_sql_list:
+                if sql_statement.strip():
+                    driver.execute_script(sql_statement)
+            execution_time = int((time.time() - start_time) * 1000)
+            if on_success:
+                on_success(execution_time)
+
         return None, execution_time
 
     def execute_downgrade(
-        self, driver: "SyncDriverAdapterBase", migration: "dict[str, Any]"
-    ) -> "tuple[Optional[str], int]":
+        self,
+        driver: "SyncDriverAdapterBase",
+        migration: "dict[str, Any]",
+        *,
+        use_transaction: "bool | None" = None,
+        on_success: "Callable[[int], None] | None" = None,
+    ) -> "tuple[str | None, int]":
         """Execute a downgrade migration.
 
         Args:
             driver: The sync database driver to use.
             migration: Migration metadata dictionary.
+            use_transaction: Override transaction behavior. If None, uses should_use_transaction logic.
+            on_success: Callback invoked with execution_time_ms before commit (for version tracking).
 
         Returns:
             Tuple of (sql_content, execution_time_ms).
@@ -278,15 +369,36 @@ class SyncMigrationRunner(BaseMigrationRunner):
         if downgrade_sql_list is None:
             return None, 0
 
+        if use_transaction is None:
+            config = self.context.config if self.context else None
+            use_transaction = self.should_use_transaction(migration, config) if config else False
+
         start_time = time.time()
 
-        for sql_statement in downgrade_sql_list:
-            if sql_statement.strip():
-                driver.execute_script(sql_statement)
-        execution_time = int((time.time() - start_time) * 1000)
+        if use_transaction:
+            try:
+                driver.begin()
+                for sql_statement in downgrade_sql_list:
+                    if sql_statement.strip():
+                        driver.execute_script(sql_statement)
+                execution_time = int((time.time() - start_time) * 1000)
+                if on_success:
+                    on_success(execution_time)
+                driver.commit()
+            except Exception:
+                driver.rollback()
+                raise
+        else:
+            for sql_statement in downgrade_sql_list:
+                if sql_statement.strip():
+                    driver.execute_script(sql_statement)
+            execution_time = int((time.time() - start_time) * 1000)
+            if on_success:
+                on_success(execution_time)
+
         return None, execution_time
 
-    def _get_migration_sql_sync(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
+    def _get_migration_sql_sync(self, migration: "dict[str, Any]", direction: str) -> "list[str] | None":
         """Get migration SQL for given direction (sync version).
 
         Args:
@@ -309,15 +421,11 @@ class SyncMigrationRunner(BaseMigrationRunner):
 
         try:
             method = loader.get_up_sql if direction == "up" else loader.get_down_sql
-            # Check if the method is async and handle appropriately
-            import inspect
-
-            if inspect.iscoroutinefunction(method):
-                # For async methods, use await_ to run in sync context
-                sql_statements = await_(method, raise_sync_error=False)(file_path)
-            else:
-                # For sync methods, call directly
-                sql_statements = method(file_path)
+            sql_statements = (
+                await_(method, raise_sync_error=False)(file_path)
+                if inspect.iscoroutinefunction(method)
+                else method(file_path)
+            )
 
         except Exception as e:
             if direction == "down":
@@ -345,11 +453,13 @@ class SyncMigrationRunner(BaseMigrationRunner):
                 for query_name in self.loader.list_queries():
                     all_queries[query_name] = self.loader.get_sql(query_name)
             else:
-                loader = get_migration_loader(file_path, self.migrations_path, self.project_root, self.context)
+                loader = get_migration_loader(
+                    file_path, self.migrations_path, self.project_root, self.context, self.loader
+                )
 
                 try:
-                    up_sql = await_(loader.get_up_sql)(file_path)
-                    down_sql = await_(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])
@@ -373,20 +483,20 @@ class AsyncMigrationRunner(BaseMigrationRunner):
         """
         return self._get_migration_files_sync()
 
-    async def load_migration(self, file_path: Path) -> "dict[str, Any]":
+    async def load_migration(self, file_path: Path, version: "str | None" = None) -> "dict[str, Any]":
         """Load a migration file and extract its components.
 
         Args:
             file_path: Path to the migration file.
+            version: Optional pre-extracted version (preserves prefixes like ext_adk_0001).
 
         Returns:
             Dictionary containing migration metadata and queries.
         """
-        # Get common metadata
-        metadata = self._load_migration_metadata_common(file_path)
+        metadata = self._load_migration_metadata_common(file_path, version)
         context_to_use = self._get_context_for_migration(file_path)
 
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use)
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, self.loader)
         loader.validate_migration_file(file_path)
 
         has_upgrade, has_downgrade = True, False
@@ -394,8 +504,6 @@ class AsyncMigrationRunner(BaseMigrationRunner):
         if file_path.suffix == ".sql":
             version = metadata["version"]
             up_query, down_query = f"migrate-{version}-up", f"migrate-{version}-down"
-            self.loader.clear_cache()
-            await async_(self.loader.load_sql)(file_path)
             has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
         else:
             try:
@@ -409,13 +517,20 @@ class AsyncMigrationRunner(BaseMigrationRunner):
         return metadata
 
     async def execute_upgrade(
-        self, driver: "AsyncDriverAdapterBase", migration: "dict[str, Any]"
-    ) -> "tuple[Optional[str], int]":
+        self,
+        driver: "AsyncDriverAdapterBase",
+        migration: "dict[str, Any]",
+        *,
+        use_transaction: "bool | None" = None,
+        on_success: "Callable[[int], Awaitable[None]] | None" = None,
+    ) -> "tuple[str | None, int]":
         """Execute an upgrade migration.
 
         Args:
             driver: The async database driver to use.
             migration: Migration metadata dictionary.
+            use_transaction: Override transaction behavior. If None, uses should_use_transaction logic.
+            on_success: Async callback invoked with execution_time_ms before commit (for version tracking).
 
         Returns:
             Tuple of (sql_content, execution_time_ms).
@@ -424,22 +539,50 @@ class AsyncMigrationRunner(BaseMigrationRunner):
         if upgrade_sql_list is None:
             return None, 0
 
+        if use_transaction is None:
+            config = self.context.config if self.context else None
+            use_transaction = self.should_use_transaction(migration, config) if config else False
+
         start_time = time.time()
 
-        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)
+        if use_transaction:
+            try:
+                await driver.begin()
+                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)
+                if on_success:
+                    await on_success(execution_time)
+                await driver.commit()
+            except Exception:
+                await driver.rollback()
+                raise
+        else:
+            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)
+            if on_success:
+                await on_success(execution_time)
+
         return None, execution_time
 
     async def execute_downgrade(
-        self, driver: "AsyncDriverAdapterBase", migration: "dict[str, Any]"
-    ) -> "tuple[Optional[str], int]":
+        self,
+        driver: "AsyncDriverAdapterBase",
+        migration: "dict[str, Any]",
+        *,
+        use_transaction: "bool | None" = None,
+        on_success: "Callable[[int], Awaitable[None]] | None" = None,
+    ) -> "tuple[str | None, int]":
         """Execute a downgrade migration.
 
         Args:
             driver: The async database driver to use.
             migration: Migration metadata dictionary.
+            use_transaction: Override transaction behavior. If None, uses should_use_transaction logic.
+            on_success: Async callback invoked with execution_time_ms before commit (for version tracking).
 
         Returns:
             Tuple of (sql_content, execution_time_ms).
@@ -448,15 +591,36 @@ class AsyncMigrationRunner(BaseMigrationRunner):
         if downgrade_sql_list is None:
             return None, 0
 
+        if use_transaction is None:
+            config = self.context.config if self.context else None
+            use_transaction = self.should_use_transaction(migration, config) if config else False
+
         start_time = time.time()
 
-        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)
+        if use_transaction:
+            try:
+                await driver.begin()
+                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)
+                if on_success:
+                    await on_success(execution_time)
+                await driver.commit()
+            except Exception:
+                await driver.rollback()
+                raise
+        else:
+            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)
+            if on_success:
+                await on_success(execution_time)
+
         return None, execution_time
 
-    async def _get_migration_sql_async(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
+    async def _get_migration_sql_async(self, migration: "dict[str, Any]", direction: str) -> "list[str] | None":
         """Get migration SQL for given direction (async version).
 
         Args:
@@ -507,7 +671,9 @@ class AsyncMigrationRunner(BaseMigrationRunner):
                 for query_name in self.loader.list_queries():
                     all_queries[query_name] = self.loader.get_sql(query_name)
             else:
-                loader = get_migration_loader(file_path, self.migrations_path, self.project_root, self.context)
+                loader = get_migration_loader(
+                    file_path, self.migrations_path, self.project_root, self.context, self.loader
+                )
 
                 try:
                     up_sql = await loader.get_up_sql(file_path)
@@ -528,7 +694,7 @@ class AsyncMigrationRunner(BaseMigrationRunner):
 def create_migration_runner(
     migrations_path: Path,
     extension_migrations: "dict[str, Path]",
-    context: "Optional[MigrationContext]",
+    context: "MigrationContext | None",
     extension_configs: "dict[str, Any]",
     is_async: "Literal[False]" = False,
 ) -> SyncMigrationRunner: ...
@@ -538,7 +704,7 @@ def create_migration_runner(
 def create_migration_runner(
     migrations_path: Path,
     extension_migrations: "dict[str, Path]",
-    context: "Optional[MigrationContext]",
+    context: "MigrationContext | None",
     extension_configs: "dict[str, Any]",
     is_async: "Literal[True]",
 ) -> AsyncMigrationRunner: ...
@@ -547,10 +713,10 @@ def create_migration_runner(
 def create_migration_runner(
     migrations_path: Path,
     extension_migrations: "dict[str, Path]",
-    context: "Optional[MigrationContext]",
+    context: "MigrationContext | None",
     extension_configs: "dict[str, Any]",
     is_async: bool = False,
-) -> "Union[SyncMigrationRunner, AsyncMigrationRunner]":
+) -> "SyncMigrationRunner | AsyncMigrationRunner":
     """Factory function to create the appropriate migration runner.
 
     Args: