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

sqlspec/migrations/base.py

@@ -4,18 +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")
 
@@ -42,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.
         """
@@ -49,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")
@@ -59,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.
@@ -89,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:
@@ -104,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
@@ -135,17 +259,31 @@ 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: "dict[str, Path] | None" = None,
+        context: "Any | None" = None,
+        extension_configs: "dict[str, dict[str, Any]] | None" = 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.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:
@@ -154,7 +292,14 @@ class BaseMigrationRunner(ABC, Generic[DriverT]):
         Returns:
             The extracted version string or None.
         """
-        parts = filename.split("_", 1)
+        from pathlib import Path
+
+        stem = Path(filename).stem
+
+        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:
@@ -172,38 +317,73 @@ 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).
         """
-        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))
-
-        return sorted(migrations, key=operator.itemgetter(0))
 
-    def _load_migration_metadata(self, file_path: Path) -> "dict[str, Any]":
+        # 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=lambda m: parse_version(m[0]))
+
+    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.
         """
-
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
+        if version is None:
+            version = self._extract_version(file_path.name)
+
+        context_to_use = self.context
+
+        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)
-        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
@@ -229,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:
@@ -292,6 +472,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 +486,56 @@ 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.
+
+        Reads extension configuration from config.extension_config for each
+        extension listed in include_extensions.
+
+        Returns:
+            Dictionary mapping extension names to their configurations.
+        """
+        configs = {}
+
+        for ext_config in self.include_extensions:
+            if not isinstance(ext_config, str):
+                logger.warning("Extension must be a string name, got: %s", ext_config)
+                continue
+
+            ext_name = ext_config
+            ext_options = getattr(self.config, "extension_config", {}).get(ext_name, {})
+            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.
@@ -320,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;
 ```
 
@@ -336,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:
@@ -363,13 +627,12 @@ 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())
 
-        (migrations_dir / ".gitkeep").touch()
-
         console.print(f"[green]Initialized migrations in {directory}[/]")
 
     @abstractmethod