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

sqlspec/migrations/base.py

@@ -4,19 +4,18 @@ 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, cast
+from typing import Any, Generic, TypeVar, cast
 
-from sqlspec._sql import sql
-from sqlspec.builder import Delete, Insert, Select
+from sqlspec.builder import Delete, Insert, Select, Update, sql
 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_
+from sqlspec.utils.version import parse_version
 
 __all__ = ("BaseMigrationCommands", "BaseMigrationRunner", "BaseMigrationTracker")
 
@@ -43,6 +42,16 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
     def _get_create_table_sql(self) -> CreateTable:
         """Get SQL builder for creating the tracking table.
 
+        Schema includes both legacy and new versioning columns:
+        - version_num: Migration version (sequential or timestamp format)
+        - version_type: Format indicator ('sequential' or 'timestamp')
+        - execution_sequence: Auto-incrementing application order
+        - description: Human-readable migration description
+        - applied_at: Timestamp when migration was applied
+        - execution_time_ms: Migration execution duration
+        - checksum: MD5 hash for content verification
+        - applied_by: User who applied the migration
+
         Returns:
             SQL builder object for table creation.
         """
@@ -50,6 +59,8 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
             sql.create_table(self.version_table)
             .if_not_exists()
             .column("version_num", "VARCHAR(32)", primary_key=True)
+            .column("version_type", "VARCHAR(16)")
+            .column("execution_sequence", "INTEGER")
             .column("description", "TEXT")
             .column("applied_at", "TIMESTAMP", default="CURRENT_TIMESTAMP", not_null=True)
             .column("execution_time_ms", "INTEGER")
@@ -60,26 +71,49 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
     def _get_current_version_sql(self) -> Select:
         """Get SQL builder for retrieving current version.
 
+        Uses execution_sequence to get the last applied migration,
+        which may differ from version_num order due to out-of-order migrations.
+
         Returns:
             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("execution_sequence DESC").limit(1)
 
     def _get_applied_migrations_sql(self) -> Select:
         """Get SQL builder for retrieving all applied migrations.
 
+        Orders by execution_sequence to show migrations in application order,
+        which preserves the actual execution history for out-of-order migrations.
+
         Returns:
             SQL builder object for migrations query.
         """
-        return sql.select("*").from_(self.version_table).order_by("version_num")
+        return sql.select("*").from_(self.version_table).order_by("execution_sequence")
+
+    def _get_next_execution_sequence_sql(self) -> Select:
+        """Get SQL builder for retrieving next execution sequence.
+
+        Returns:
+            SQL builder object for sequence query.
+        """
+        return sql.select("COALESCE(MAX(execution_sequence), 0) + 1 AS next_seq").from_(self.version_table)
 
     def _get_record_migration_sql(
-        self, version: str, description: str, execution_time_ms: int, checksum: str, applied_by: str
+        self,
+        version: str,
+        version_type: str,
+        execution_sequence: int,
+        description: str,
+        execution_time_ms: int,
+        checksum: str,
+        applied_by: str,
     ) -> Insert:
         """Get SQL builder for recording a migration.
 
         Args:
             version: Version number of the migration.
+            version_type: Version format type ('sequential' or 'timestamp').
+            execution_sequence: Auto-incrementing application order.
             description: Description of the migration.
             execution_time_ms: Execution time in milliseconds.
             checksum: MD5 checksum of the migration content.
@@ -90,8 +124,16 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
         """
         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)
+            .columns(
+                "version_num",
+                "version_type",
+                "execution_sequence",
+                "description",
+                "execution_time_ms",
+                "checksum",
+                "applied_by",
+            )
+            .values(version, version_type, execution_sequence, description, execution_time_ms, checksum, applied_by)
         )
 
     def _get_remove_migration_sql(self, version: str) -> Delete:
@@ -105,9 +147,90 @@ class BaseMigrationTracker(ABC, Generic[DriverT]):
         """
         return sql.delete().from_(self.version_table).where(sql.version_num == version)
 
+    def _get_update_version_sql(self, old_version: str, new_version: str, new_version_type: str) -> Update:
+        """Get SQL builder for updating version record.
+
+        Updates version_num and version_type while preserving execution_sequence,
+        applied_at, and other metadata. Used during fix command to convert
+        timestamp versions to sequential format.
+
+        Args:
+            old_version: Current version string.
+            new_version: New version string.
+            new_version_type: New version type ('sequential' or 'timestamp').
+
+        Returns:
+            SQL builder object for update.
+        """
+        return (
+            sql.update(self.version_table)
+            .set("version_num", new_version)
+            .set("version_type", new_version_type)
+            .where(sql.version_num == old_version)
+        )
+
+    def _get_check_column_exists_sql(self) -> Select:
+        """Get SQL to check what columns exist in the tracking table.
+
+        Returns a query that will fail gracefully if the table doesn't exist,
+        and returns column names if it does.
+
+        Returns:
+            SQL builder object for column check query.
+        """
+        return sql.select("*").from_(self.version_table).limit(0)
+
+    def _get_add_missing_columns_sql(self, missing_columns: "set[str]") -> "list[str]":
+        """Generate ALTER TABLE statements to add missing columns.
+
+        Args:
+            missing_columns: Set of column names that need to be added.
+
+        Returns:
+            List of SQL statements to execute.
+        """
+
+        statements = []
+        target_create = self._get_create_table_sql()
+
+        column_definitions = {col.name.lower(): col for col in target_create.columns}
+
+        for col_name in sorted(missing_columns):
+            if col_name in column_definitions:
+                col_def = column_definitions[col_name]
+                alter = sql.alter_table(self.version_table).add_column(
+                    name=col_def.name,
+                    dtype=col_def.dtype,
+                    default=col_def.default,
+                    not_null=col_def.not_null,
+                    unique=col_def.unique,
+                    comment=col_def.comment,
+                )
+                statements.append(str(alter))
+
+        return statements
+
+    def _detect_missing_columns(self, existing_columns: "set[str]") -> "set[str]":
+        """Detect which columns are missing from the current schema.
+
+        Args:
+            existing_columns: Set of existing column names (may be uppercase/lowercase).
+
+        Returns:
+            Set of missing column names (lowercase).
+        """
+        target_create = self._get_create_table_sql()
+        target_columns = {col.name.lower() for col in target_create.columns}
+        existing_lower = {col.lower() for col in existing_columns}
+        return target_columns - existing_lower
+
     @abstractmethod
     def ensure_tracking_table(self, driver: DriverT) -> Any:
-        """Create the migration tracking table if it doesn't exist."""
+        """Create the migration tracking table if it doesn't exist.
+
+        Implementations should also check for and add any missing columns
+        to support schema migrations from older versions.
+        """
         ...
 
     @abstractmethod
@@ -141,9 +264,9 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
     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,
+        extension_migrations: "dict[str, Path] | None" = None,
+        context: "Any | None" = None,
+        extension_configs: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize the migration runner.
 
@@ -156,11 +279,11 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         self.migrations_path = migrations_path
         self.extension_migrations = extension_migrations or {}
         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.
 
         Args:
@@ -169,13 +292,14 @@ 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
+        from pathlib import Path
+
+        stem = Path(filename).stem
 
-        # Regular version extraction
-        parts = filename.split("_", 1)
+        if stem.startswith("ext_"):
+            return stem
+
+        parts = stem.split("_", 1)
         return parts[0].zfill(4) if parts and parts[0].isdigit() else None
 
     def _calculate_checksum(self, content: str) -> str:
@@ -193,6 +317,9 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
     def _get_migration_files_sync(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
 
+        Uses version-aware sorting that handles both sequential and timestamp
+        formats correctly, with extension migrations sorted by extension name.
+
         Returns:
             List of tuples containing (version, file_path).
         """
@@ -222,42 +349,23 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
                             prefixed_version = f"ext_{ext_name}_{version}"
                             migrations.append((prefixed_version, file_path))
 
-        return sorted(migrations, key=operator.itemgetter(0))
+        return sorted(migrations, key=lambda m: parse_version(m[0]))
 
-    def _load_migration_metadata(self, file_path: Path) -> "dict[str, Any]":
+    def _load_migration_metadata(self, file_path: Path, version: "str | None" = None) -> "dict[str, Any]":
         """Load migration metadata from file.
 
         Args:
             file_path: Path to the migration file.
+            version: Optional pre-extracted version (preserves prefixes like ext_adk_0001).
 
         Returns:
             Migration metadata dictionary.
         """
+        if version is None:
+            version = self._extract_version(file_path.name)
 
-        # 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:
@@ -276,7 +384,6 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         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
@@ -302,7 +409,7 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
             "loader": loader,
         }
 
-    def _get_migration_sql(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
+    def _get_migration_sql(self, migration: "dict[str, Any]", direction: str) -> "list[str] | None":
         """Get migration SQL for given direction.
 
         Args:
@@ -385,8 +492,8 @@ class BaseMigrationCommands(ABC, Generic[ConfigT, DriverT]):
     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).
+        Reads extension configuration from config.extension_config for each
+        extension listed in include_extensions.
 
         Returns:
             Dictionary mapping extension names to their configurations.
@@ -394,28 +501,12 @@ class BaseMigrationCommands(ABC, Generic[ConfigT, DriverT]):
         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)
+            if not isinstance(ext_config, str):
+                logger.warning("Extension must be a string name, got: %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"
-
+            ext_name = ext_config
+            ext_options = getattr(self.config, "extension_config", {}).get(ext_name, {})
             configs[ext_name] = ext_options
 
         return configs
@@ -461,13 +552,13 @@ This directory contains database migration files.
 Migration files use SQLFileLoader's named query syntax with versioned names:
 
 ```sql
--- name: migrate-0001-up
+-- name: migrate-20251011120000-up
 CREATE TABLE example (
     id INTEGER PRIMARY KEY,
     name TEXT NOT NULL
 );
 
--- name: migrate-0001-down
+-- name: migrate-20251011120000-down
 DROP TABLE example;
 ```
 
@@ -477,16 +568,48 @@ DROP TABLE example;
 
 Format: `{version}_{description}.sql`
 
-- Version: Zero-padded 4-digit number (0001, 0002, etc.)
+- Version: Timestamp in YYYYMMDDHHmmss format (UTC)
 - Description: Brief description using underscores
-- Example: `0001_create_users_table.sql`
+- Example: `20251011120000_create_users_table.sql`
 
 ### Query Names
 
 - Upgrade: `migrate-{version}-up`
 - Downgrade: `migrate-{version}-down`
 
-This naming ensures proper sorting and avoids conflicts when loading multiple files.
+## Version Format
+
+Migrations use **timestamp-based versioning** (YYYYMMDDHHmmss):
+
+- **Format**: 14-digit UTC timestamp
+- **Example**: `20251011120000` (October 11, 2025 at 12:00:00 UTC)
+- **Benefits**: Eliminates merge conflicts when multiple developers create migrations concurrently
+
+### Creating Migrations
+
+Use the CLI to generate timestamped migrations:
+
+```bash
+sqlspec create-migration "add user table"
+# Creates: 20251011120000_add_user_table.sql
+```
+
+The timestamp is automatically generated in UTC timezone.
+
+## Migration Execution
+
+Migrations are applied in chronological order based on their timestamps.
+The database tracks both version and execution order separately to handle
+out-of-order migrations gracefully (e.g., from late-merging branches).
+"""
+
+    def _get_init_init_content(self) -> str:
+        """Get __init__.py content for migration directory initialization.
+
+        Returns:
+            Python module docstring content for the __init__.py file.
+        """
+        return """Migrations.
 """
 
     def init_directory(self, directory: str, package: bool = True) -> None:
@@ -504,7 +627,8 @@ This naming ensures proper sorting and avoids conflicts when loading multiple fi
         migrations_dir.mkdir(parents=True, exist_ok=True)
 
         if package:
-            (migrations_dir / "__init__.py").touch()
+            init = migrations_dir / "__init__.py"
+            init.write_text(self._get_init_init_content())
 
         readme = migrations_dir / "README.md"
         readme.write_text(self._get_init_readme_content())