sqlspec 0.13.1__py3-none-any.whl → 0.14.1__py3-none-any.whl

sqlspec/migrations/runner.py

@@ -0,0 +1,215 @@
+"""Migration execution engine for SQLSpec.
+
+This module handles migration file loading and execution using SQLFileLoader.
+"""
+
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from sqlspec.migrations.base import BaseMigrationRunner
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
+    from sqlspec.statement.sql import SQL
+
+__all__ = ("AsyncMigrationRunner", "SyncMigrationRunner")
+
+logger = get_logger("migrations.runner")
+
+
+class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterProtocol[Any]"]):
+    """Sync version - executes migrations using SQLFileLoader."""
+
+    def get_migration_files(self) -> "list[tuple[str, Path]]":
+        """Get all migration files sorted by version.
+
+        Returns:
+            List of (version, path) tuples sorted by version.
+        """
+        return self._get_migration_files_sync()
+
+    def load_migration(self, file_path: Path) -> "dict[str, Any]":
+        """Load a migration file and extract its components.
+
+        Args:
+            file_path: Path to the migration file.
+
+        Returns:
+            Dictionary containing migration metadata and queries.
+        """
+        return self._load_migration_metadata(file_path)
+
+    def execute_upgrade(
+        self, driver: "SyncDriverAdapterProtocol[Any]", migration: "dict[str, Any]"
+    ) -> "tuple[Optional[str], int]":
+        """Execute an upgrade migration.
+
+        Args:
+            driver: The database driver to use.
+            migration: Migration metadata dictionary.
+
+        Returns:
+            Tuple of (sql_content, execution_time_ms).
+        """
+        upgrade_sql = self._get_migration_sql(migration, "up")
+        if upgrade_sql is None:
+            return None, 0
+
+        start_time = time.time()
+
+        # Execute migration
+        driver.execute(upgrade_sql)
+
+        execution_time = int((time.time() - start_time) * 1000)
+        return None, execution_time
+
+    def execute_downgrade(
+        self, driver: "SyncDriverAdapterProtocol[Any]", migration: "dict[str, Any]"
+    ) -> "tuple[Optional[str], int]":
+        """Execute a downgrade migration.
+
+        Args:
+            driver: The database driver to use.
+            migration: Migration metadata dictionary.
+
+        Returns:
+            Tuple of (sql_content, execution_time_ms).
+        """
+        downgrade_sql = self._get_migration_sql(migration, "down")
+        if downgrade_sql is None:
+            return None, 0
+
+        start_time = time.time()
+
+        # Execute migration
+        driver.execute(downgrade_sql)
+
+        execution_time = int((time.time() - start_time) * 1000)
+        return None, execution_time
+
+    def load_all_migrations(self) -> "dict[str, SQL]":
+        """Load all migrations into a single namespace for bulk operations.
+
+        Returns a dictionary mapping query names to SQL objects.
+        Useful for:
+        - Migration analysis tools
+        - Documentation generation
+        - Validation and linting
+        - Migration squashing
+
+        Returns:
+            Dictionary mapping query names to SQL objects.
+        """
+        all_queries = {}
+        migrations = self.get_migration_files()
+
+        for _version, file_path in migrations:
+            self.loader.load_sql(file_path)
+
+            # Get all queries from this file
+            for query_name in self.loader.list_queries():
+                # Store with full query name for uniqueness
+                all_queries[query_name] = self.loader.get_sql(query_name)
+
+        return all_queries
+
+
+class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterProtocol[Any]"]):
+    """Async version - executes migrations using SQLFileLoader."""
+
+    async def get_migration_files(self) -> "list[tuple[str, Path]]":
+        """Get all migration files sorted by version.
+
+        Returns:
+            List of tuples containing (version, file_path).
+        """
+        # For async, we still use the sync file operations since Path.glob is sync
+        return self._get_migration_files_sync()
+
+    async def load_migration(self, file_path: Path) -> "dict[str, Any]":
+        """Load a migration file and extract its components.
+
+        Args:
+            file_path: Path to the migration file.
+
+        Returns:
+            Dictionary containing migration metadata.
+        """
+        # File loading is still sync, so we use the base implementation
+        return self._load_migration_metadata(file_path)
+
+    async def execute_upgrade(
+        self, driver: "AsyncDriverAdapterProtocol[Any]", migration: "dict[str, Any]"
+    ) -> "tuple[Optional[str], int]":
+        """Execute an upgrade migration.
+
+        Args:
+            driver: The async database driver to use.
+            migration: Migration metadata dictionary.
+
+        Returns:
+            Tuple of (sql_content, execution_time_ms).
+        """
+        upgrade_sql = self._get_migration_sql(migration, "up")
+        if upgrade_sql is None:
+            return None, 0
+
+        start_time = time.time()
+
+        # Execute migration
+        await driver.execute(upgrade_sql)
+
+        execution_time = int((time.time() - start_time) * 1000)
+        return None, execution_time
+
+    async def execute_downgrade(
+        self, driver: "AsyncDriverAdapterProtocol[Any]", migration: "dict[str, Any]"
+    ) -> "tuple[Optional[str], int]":
+        """Execute a downgrade migration.
+
+        Args:
+            driver: The async database driver to use.
+            migration: Migration metadata dictionary.
+
+        Returns:
+            Tuple of (sql_content, execution_time_ms).
+        """
+        downgrade_sql = self._get_migration_sql(migration, "down")
+        if downgrade_sql is None:
+            return None, 0
+
+        start_time = time.time()
+
+        # Execute migration
+        await driver.execute(downgrade_sql)
+
+        execution_time = int((time.time() - start_time) * 1000)
+        return None, execution_time
+
+    async def load_all_migrations(self) -> "dict[str, SQL]":
+        """Load all migrations into a single namespace for bulk operations.
+
+        Returns a dictionary mapping query names to SQL objects.
+        Useful for:
+        - Migration analysis tools
+        - Documentation generation
+        - Validation and linting
+        - Migration squashing
+
+        Returns:
+            Dictionary mapping query names to SQL objects.
+        """
+        all_queries = {}
+        migrations = await self.get_migration_files()
+
+        for _version, file_path in migrations:
+            self.loader.load_sql(file_path)
+
+            # Get all queries from this file
+            for query_name in self.loader.list_queries():
+                # Store with full query name for uniqueness
+                all_queries[query_name] = self.loader.get_sql(query_name)
+
+        return all_queries

sqlspec/migrations/tracker.py

@@ -0,0 +1,153 @@
+"""Migration version tracking for SQLSpec.
+
+This module provides functionality to track applied migrations in the database.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from sqlspec.migrations.base import BaseMigrationTracker
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
+
+__all__ = ("AsyncMigrationTracker", "SyncMigrationTracker")
+
+
+class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"]):
+    """Sync version - tracks applied migrations in the database."""
+
+    def ensure_tracking_table(self, driver: "SyncDriverAdapterProtocol[Any]") -> None:
+        """Create the migration tracking table if it doesn't exist.
+
+        Args:
+            driver: The database driver to use.
+        """
+        driver.execute(self._get_create_table_sql())
+
+    def get_current_version(self, driver: "SyncDriverAdapterProtocol[Any]") -> Optional[str]:
+        """Get the latest applied migration version.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            The current version number or None if no migrations applied.
+        """
+        result = driver.execute(self._get_current_version_sql())
+        return result.data[0]["version_num"] if result.data else None
+
+    def get_applied_migrations(self, driver: "SyncDriverAdapterProtocol[Any]") -> "list[dict[str, Any]]":
+        """Get all applied migrations in order.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            List of migration records.
+        """
+        result = driver.execute(self._get_applied_migrations_sql())
+        return result.data
+
+    def record_migration(
+        self,
+        driver: "SyncDriverAdapterProtocol[Any]",
+        version: str,
+        description: str,
+        execution_time_ms: int,
+        checksum: str,
+    ) -> None:
+        """Record a successfully applied migration.
+
+        Args:
+            driver: The database driver to use.
+            version: Version number of the migration.
+            description: Description of the migration.
+            execution_time_ms: Execution time in milliseconds.
+            checksum: MD5 checksum of the migration content.
+            connection: Optional connection to use for the operation.
+        """
+        import os
+
+        applied_by = os.environ.get("USER", "unknown")
+
+        driver.execute(self._get_record_migration_sql(version, description, execution_time_ms, checksum, applied_by))
+
+    def remove_migration(self, driver: "SyncDriverAdapterProtocol[Any]", version: str) -> None:
+        """Remove a migration record (used during downgrade).
+
+        Args:
+            driver: The database driver to use.
+            version: Version number to remove.
+            connection: Optional connection to use for the operation.
+        """
+        driver.execute(self._get_remove_migration_sql(version))
+
+
+class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterProtocol[Any]"]):
+    """Async version - tracks applied migrations in the database."""
+
+    async def ensure_tracking_table(self, driver: "AsyncDriverAdapterProtocol[Any]") -> None:
+        """Create the migration tracking table if it doesn't exist.
+
+        Args:
+            driver: The async database driver to use.
+        """
+        await driver.execute(self._get_create_table_sql())
+
+    async def get_current_version(self, driver: "AsyncDriverAdapterProtocol[Any]") -> Optional[str]:
+        """Get the latest applied migration version.
+
+        Args:
+            driver: The async database driver to use.
+
+        Returns:
+            The current version number or None if no migrations applied.
+        """
+        result = await driver.execute(self._get_current_version_sql())
+        return result.data[0]["version_num"] if result.data else None
+
+    async def get_applied_migrations(self, driver: "AsyncDriverAdapterProtocol[Any]") -> "list[dict[str, Any]]":
+        """Get all applied migrations in order.
+
+        Args:
+            driver: The async database driver to use.
+
+        Returns:
+            List of migration records.
+        """
+        result = await driver.execute(self._get_applied_migrations_sql())
+        return result.data
+
+    async def record_migration(
+        self,
+        driver: "AsyncDriverAdapterProtocol[Any]",
+        version: str,
+        description: str,
+        execution_time_ms: int,
+        checksum: str,
+    ) -> None:
+        """Record a successfully applied migration.
+
+        Args:
+            driver: The async database driver to use.
+            version: Version number of the migration.
+            description: Description of the migration.
+            execution_time_ms: Execution time in milliseconds.
+            checksum: MD5 checksum of the migration content.
+        """
+        import os
+
+        applied_by = os.environ.get("USER", "unknown")
+
+        await driver.execute(
+            self._get_record_migration_sql(version, description, execution_time_ms, checksum, applied_by)
+        )
+
+    async def remove_migration(self, driver: "AsyncDriverAdapterProtocol[Any]", version: str) -> None:
+        """Remove a migration record (used during downgrade).
+
+        Args:
+            driver: The async database driver to use.
+            version: Version number to remove.
+        """
+        await driver.execute(self._get_remove_migration_sql(version))

sqlspec/migrations/utils.py

@@ -0,0 +1,89 @@
+"""Utility functions for SQLSpec migrations.
+
+This module provides helper functions for migration operations.
+"""
+
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterProtocol
+
+__all__ = ("create_migration_file", "drop_all", "get_author")
+
+
+def create_migration_file(migrations_dir: Path, version: str, message: str) -> Path:
+    """Create a new migration file from template.
+
+    Args:
+        migrations_dir: Directory to create the migration in.
+        version: Version number for the migration.
+        message: Description message for the migration.
+
+    Returns:
+        Path to the created migration file.
+    """
+    # Sanitize message for filename
+    safe_message = message.lower()
+    safe_message = "".join(c if c.isalnum() or c in " -" else "" for c in safe_message)
+    safe_message = safe_message.replace(" ", "_").replace("-", "_")
+    safe_message = "_".join(filter(None, safe_message.split("_")))[:50]
+
+    filename = f"{version}_{safe_message}.sql"
+    file_path = migrations_dir / filename
+
+    # Generate template content
+    template = f"""-- SQLSpec Migration
+-- Version: {version}
+-- Description: {message}
+-- Created: {datetime.now(timezone.utc).isoformat()}
+-- Author: {get_author()}
+
+-- name: migrate-{version}-up
+-- TODO: Add your upgrade SQL statements here
+-- Example:
+-- CREATE TABLE example (
+--     id INTEGER PRIMARY KEY,
+--     name TEXT NOT NULL
+-- );
+
+-- name: migrate-{version}-down
+-- TODO: Add your downgrade SQL statements here (optional)
+-- Example:
+-- DROP TABLE example;
+"""
+
+    file_path.write_text(template)
+    return file_path
+
+
+def get_author() -> str:
+    """Get current user for migration metadata.
+
+    Returns:
+        Username from environment or 'unknown'.
+    """
+    return os.environ.get("USER", "unknown")
+
+
+async def drop_all(
+    engine: "AsyncDriverAdapterProtocol[Any]", version_table_name: str, metadata: Optional[Any] = None
+) -> None:
+    """Drop all tables from the database.
+
+    This is a placeholder for database-specific implementations.
+
+    Args:
+        engine: The database engine/driver.
+        version_table_name: Name of the version tracking table.
+        metadata: Optional metadata object.
+
+    Raises:
+        NotImplementedError: Always, as this requires database-specific logic.
+    """
+    # This would need database-specific implementation
+    # For now, it's a placeholder
+    msg = "drop_all functionality requires database-specific implementation"
+    raise NotImplementedError(msg)

sqlspec/protocols.py

@@ -30,12 +30,16 @@ __all__ = (
     "DictProtocol",
     "FilterAppenderProtocol",
     "FilterParameterProtocol",
+    "HasExpressionProtocol",
     "HasExpressionsProtocol",
     "HasLimitProtocol",
     "HasOffsetProtocol",
     "HasOrderByProtocol",
+    "HasParameterBuilderProtocol",
     "HasRiskLevelProtocol",
+    "HasSQLGlotExpressionProtocol",
     "HasSQLMethodProtocol",
+    "HasToStatementProtocol",
     "HasWhereProtocol",
     "IndexableRow",
     "IterableParameters",
@@ -501,9 +505,39 @@ class ObjectStoreProtocol(Protocol):
         raise NotImplementedError(msg)
 
 
-# =============================================================================
-# SQL Builder Protocols
-# =============================================================================
+@runtime_checkable
+class HasSQLGlotExpressionProtocol(Protocol):
+    """Protocol for objects with a sqlglot_expression property."""
+
+    @property
+    def sqlglot_expression(self) -> "Optional[exp.Expression]":
+        """Return the SQLGlot expression for this object."""
+        ...
+
+
+@runtime_checkable
+class HasParameterBuilderProtocol(Protocol):
+    """Protocol for objects that can add parameters."""
+
+    def add_parameter(self, value: Any, name: "Optional[str]" = None) -> tuple[Any, str]:
+        """Add a parameter to the builder."""
+        ...
+
+
+@runtime_checkable
+class HasExpressionProtocol(Protocol):
+    """Protocol for objects with an _expression attribute."""
+
+    _expression: "Optional[exp.Expression]"
+
+
+@runtime_checkable
+class HasToStatementProtocol(Protocol):
+    """Protocol for objects with a to_statement method."""
+
+    def to_statement(self) -> Any:
+        """Convert to SQL statement."""
+        ...
 
 
 @runtime_checkable

sqlspec/statement/builder/__init__.py

@@ -7,9 +7,9 @@ parameter binding and validation.
 """
 
 from sqlspec.exceptions import SQLBuilderError
-from sqlspec.statement.builder.base import QueryBuilder, SafeQuery
-from sqlspec.statement.builder.column import Column, ColumnExpression, FunctionColumn
-from sqlspec.statement.builder.ddl import (
+from sqlspec.statement.builder._base import QueryBuilder, SafeQuery
+from sqlspec.statement.builder._column import Column, ColumnExpression, FunctionColumn
+from sqlspec.statement.builder._ddl import (
     AlterTable,
     CommentOn,
     CreateIndex,
@@ -26,12 +26,12 @@ from sqlspec.statement.builder.ddl import (
     RenameTable,
     TruncateTable,
 )
-from sqlspec.statement.builder.delete import Delete
-from sqlspec.statement.builder.insert import Insert
-from sqlspec.statement.builder.merge import Merge
+from sqlspec.statement.builder._delete import Delete
+from sqlspec.statement.builder._insert import Insert
+from sqlspec.statement.builder._merge import Merge
+from sqlspec.statement.builder._select import Select
+from sqlspec.statement.builder._update import Update
 from sqlspec.statement.builder.mixins import WhereClauseMixin
-from sqlspec.statement.builder.select import Select
-from sqlspec.statement.builder.update import Update
 
 __all__ = (
     "AlterTable",