sqlspec 0.14.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

sqlspec/migrations/commands.py

@@ -8,11 +8,11 @@ from typing import TYPE_CHECKING, Any, Union, cast
 from rich.console import Console
 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.statement.sql import SQL
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.sync_tools import await_
 
@@ -26,7 +26,7 @@ console = Console()
 
 
 class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
-    """SQLSpec native migration commands (sync version)."""
+    """SQLSpec native migration commands."""
 
     def __init__(self, config: "SyncConfigT") -> None:
         """Initialize migration commands.
@@ -95,8 +95,6 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
 
             current = self.tracker.get_current_version(driver)
             all_migrations = self.runner.get_migration_files()
-
-            # Determine pending migrations
             pending = []
             for version, file_path in all_migrations:
                 if (current is None or version > current) and (revision == "head" or version <= revision):
@@ -108,21 +106,16 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
 
             console.print(f"[yellow]Found {len(pending)} pending migrations[/]")
 
-            # Execute migrations
             for version, file_path in pending:
                 migration = self.runner.load_migration(file_path)
 
                 console.print(f"\n[cyan]Applying {version}:[/] {migration['description']}")
 
                 try:
-                    # Execute migration
                     _, execution_time = self.runner.execute_upgrade(driver, migration)
-
-                    # Record in tracking table
                     self.tracker.record_migration(
                         driver, migration["version"], migration["description"], execution_time, migration["checksum"]
                     )
-
                     console.print(f"[green]✓ Applied in {execution_time}ms[/]")
 
                 except Exception as e:
@@ -137,19 +130,14 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
         """
         with self.config.provide_session() as driver:
             self.tracker.ensure_tracking_table(driver)
-
             applied = self.tracker.get_applied_migrations(driver)
             if not applied:
                 console.print("[yellow]No migrations to downgrade[/]")
                 return
-
-            # Determine migrations to revert
             to_revert = []
             if revision == "-1":
-                # Downgrade one step
                 to_revert = [applied[-1]]
             else:
-                # Downgrade to specific version
                 for migration in reversed(applied):
                     if migration["version_num"] > revision:
                         to_revert.append(migration)
@@ -159,29 +147,18 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
                 return
 
             console.print(f"[yellow]Reverting {len(to_revert)} migrations[/]")
-
-            # Load migration files
             all_files = dict(self.runner.get_migration_files())
-
             for migration_record in to_revert:
                 version = migration_record["version_num"]
-
                 if version not in all_files:
                     console.print(f"[red]Migration file not found for {version}[/]")
                     continue
-
                 migration = self.runner.load_migration(all_files[version])
                 console.print(f"\n[cyan]Reverting {version}:[/] {migration['description']}")
-
                 try:
-                    # Execute downgrade
                     _, execution_time = self.runner.execute_downgrade(driver, migration)
-
-                    # Remove from tracking table
                     self.tracker.remove_migration(driver, version)
-
                     console.print(f"[green]✓ Reverted in {execution_time}ms[/]")
-
                 except Exception as e:
                     console.print(f"[red]✗ Failed: {e}[/]")
                     raise
@@ -194,51 +171,37 @@ class SyncMigrationCommands(BaseMigrationCommands["SyncConfigT", Any]):
         """
         with self.config.provide_session() as driver:
             self.tracker.ensure_tracking_table(driver)
-
-            # Validate revision exists
             all_migrations = dict(self.runner.get_migration_files())
             if revision not in all_migrations:
                 console.print(f"[red]Unknown revision: {revision}[/]")
                 return
-
-            # Clear existing records and stamp
-            clear_sql = SQL(f"DELETE FROM {self.tracker.version_table}")
+            clear_sql = sql.delete().from_(self.tracker.version_table)
             driver.execute(clear_sql)
-
             self.tracker.record_migration(driver, revision, f"Stamped to {revision}", 0, "manual-stamp")
-
             console.print(f"[green]Database stamped at revision {revision}[/]")
 
-    def revision(self, message: str) -> None:
+    def revision(self, message: str, file_type: str = "sql") -> None:
         """Create a new migration file.
 
         Args:
             message: Description for the migration.
+            file_type: Type of migration file to create ('sql' or 'py').
         """
-        # Determine next version number
         existing = self.runner.get_migration_files()
-        if existing:
-            last_version = existing[-1][0]
-            next_num = int(last_version) + 1
-        else:
-            next_num = 1
-
+        next_num = int(existing[-1][0]) + 1 if existing else 1
         next_version = str(next_num).zfill(4)
-
-        # Create migration file
-        file_path = create_migration_file(self.migrations_path, next_version, message)
-
+        file_path = create_migration_file(self.migrations_path, next_version, message, file_type)
         console.print(f"[green]Created migration:[/] {file_path}")
 
 
 class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
-    """SQLSpec native async migration commands."""
+    """SQLSpec native migration commands."""
 
     def __init__(self, sqlspec_config: "AsyncConfigT") -> None:
-        """Initialize async migration commands.
+        """Initialize migration commands.
 
         Args:
-            sqlspec_config: The async SQLSpec configuration.
+            sqlspec_config: The SQLSpec configuration.
         """
         super().__init__(sqlspec_config)
         self.tracker = AsyncMigrationTracker(self.version_table)
@@ -251,7 +214,6 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
             directory: Directory path for migrations.
             package: Whether to create __init__.py in the directory.
         """
-        # For async, we still use sync directory operations
         self.init_directory(directory, package)
 
     async def current(self, verbose: bool = False) -> None:
@@ -269,17 +231,14 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
                 return
 
             console.print(f"[green]Current version:[/] {current}")
-
             if verbose:
                 applied = await self.tracker.get_applied_migrations(driver)
-
                 table = Table(title="Applied Migrations")
                 table.add_column("Version", style="cyan")
                 table.add_column("Description")
                 table.add_column("Applied At")
                 table.add_column("Time (ms)", justify="right")
                 table.add_column("Applied By")
-
                 for migration in applied:
                     table.add_row(
                         migration["version_num"],
@@ -288,50 +247,36 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
                         str(migration.get("execution_time_ms", "")),
                         migration.get("applied_by", ""),
                     )
-
                 console.print(table)
 
     async def upgrade(self, revision: str = "head") -> None:
         """Upgrade to a target revision.
 
         Args:
-            revision: Target revision (default: "head" for latest).
+            revision: Target revision or "head" for latest.
         """
         async with self.config.provide_session() as driver:
             await self.tracker.ensure_tracking_table(driver)
 
             current = await self.tracker.get_current_version(driver)
             all_migrations = await self.runner.get_migration_files()
-
-            # Determine pending migrations
             pending = []
             for version, file_path in all_migrations:
                 if (current is None or version > current) and (revision == "head" or version <= revision):
                     pending.append((version, file_path))
-
             if not pending:
                 console.print("[green]Already at latest version[/]")
                 return
-
             console.print(f"[yellow]Found {len(pending)} pending migrations[/]")
-
-            # Execute migrations
             for version, file_path in pending:
                 migration = await self.runner.load_migration(file_path)
-
                 console.print(f"\n[cyan]Applying {version}:[/] {migration['description']}")
-
                 try:
-                    # Execute migration
                     _, execution_time = await self.runner.execute_upgrade(driver, migration)
-
-                    # Record in tracking table
                     await self.tracker.record_migration(
                         driver, migration["version"], migration["description"], execution_time, migration["checksum"]
                     )
-
                     console.print(f"[green]✓ Applied in {execution_time}ms[/]")
-
                 except Exception as e:
                     console.print(f"[red]✗ Failed: {e}[/]")
                     raise
@@ -340,7 +285,7 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
         """Downgrade to a target revision.
 
         Args:
-            revision: Target revision (default: "-1" for one step back).
+            revision: Target revision or "-1" for one step back.
         """
         async with self.config.provide_session() as driver:
             await self.tracker.ensure_tracking_table(driver)
@@ -349,30 +294,21 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
             if not applied:
                 console.print("[yellow]No migrations to downgrade[/]")
                 return
-
-            # Determine migrations to revert
             to_revert = []
             if revision == "-1":
-                # Downgrade one step
                 to_revert = [applied[-1]]
             else:
-                # Downgrade to specific version
                 for migration in reversed(applied):
                     if migration["version_num"] > revision:
                         to_revert.append(migration)
-
             if not to_revert:
                 console.print("[yellow]Nothing to downgrade[/]")
                 return
 
             console.print(f"[yellow]Reverting {len(to_revert)} migrations[/]")
-
-            # Load migration files
             all_files = dict(await self.runner.get_migration_files())
-
             for migration_record in to_revert:
                 version = migration_record["version_num"]
-
                 if version not in all_files:
                     console.print(f"[red]Migration file not found for {version}[/]")
                     continue
@@ -381,14 +317,9 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
                 console.print(f"\n[cyan]Reverting {version}:[/] {migration['description']}")
 
                 try:
-                    # Execute downgrade
                     _, execution_time = await self.runner.execute_downgrade(driver, migration)
-
-                    # Remove from tracking table
                     await self.tracker.remove_migration(driver, version)
-
                     console.print(f"[green]✓ Reverted in {execution_time}ms[/]")
-
                 except Exception as e:
                     console.print(f"[red]✗ Failed: {e}[/]")
                     raise
@@ -402,39 +333,27 @@ class AsyncMigrationCommands(BaseMigrationCommands["AsyncConfigT", Any]):
         async with self.config.provide_session() as driver:
             await self.tracker.ensure_tracking_table(driver)
 
-            # Validate revision exists
             all_migrations = dict(await self.runner.get_migration_files())
             if revision not in all_migrations:
                 console.print(f"[red]Unknown revision: {revision}[/]")
                 return
 
-            # Clear existing records and stamp
-            clear_sql = SQL(f"DELETE FROM {self.tracker.version_table}")
+            clear_sql = sql.delete().from_(self.tracker.version_table)
             await driver.execute(clear_sql)
-
             await self.tracker.record_migration(driver, revision, f"Stamped to {revision}", 0, "manual-stamp")
-
             console.print(f"[green]Database stamped at revision {revision}[/]")
 
-    async def revision(self, message: str) -> None:
+    async def revision(self, message: str, file_type: str = "sql") -> None:
         """Create a new migration file.
 
         Args:
-            message: Description of the migration.
+            message: Description for the migration.
+            file_type: Type of migration file to create ('sql' or 'py').
         """
-        # Determine next version number
         existing = await self.runner.get_migration_files()
-        if existing:
-            last_version = existing[-1][0]
-            next_num = int(last_version) + 1
-        else:
-            next_num = 1
-
+        next_num = int(existing[-1][0]) + 1 if existing else 1
         next_version = str(next_num).zfill(4)
-
-        # Create migration file
-        file_path = create_migration_file(self.migrations_path, next_version, message)
-
+        file_path = create_migration_file(self.migrations_path, next_version, message, file_type)
         console.print(f"[green]Created migration:[/] {file_path}")
 
 
@@ -442,19 +361,17 @@ class MigrationCommands:
     """Unified migration commands that adapt to sync/async configs."""
 
     def __init__(self, config: "Union[SyncConfigT, AsyncConfigT]") -> None:
-        """Initialize migration commands with appropriate sync/async implementation.
+        """Initialize migration commands with sync/async implementation.
 
         Args:
-            config: The SQLSpec configuration (sync or async).
+            config: The SQLSpec configuration.
         """
-
         if config.is_async:
             self._impl: Union[AsyncMigrationCommands[Any], SyncMigrationCommands[Any]] = AsyncMigrationCommands(
                 cast("AsyncConfigT", config)
             )
         else:
             self._impl = SyncMigrationCommands(cast("SyncConfigT", config))
-
         self._is_async = config.is_async
 
     def init(self, directory: str, package: bool = True) -> None:
@@ -513,13 +430,14 @@ class MigrationCommands:
         else:
             cast("SyncMigrationCommands[Any]", self._impl).stamp(revision)
 
-    def revision(self, message: str) -> None:
+    def revision(self, message: str, file_type: str = "sql") -> None:
         """Create a new migration file.
 
         Args:
             message: Description for the migration.
+            file_type: Type of migration file to create ('sql' or 'py').
         """
         if self._is_async:
-            await_(cast("AsyncMigrationCommands[Any]", self._impl).revision, raise_sync_error=False)(message)
+            await_(cast("AsyncMigrationCommands[Any]", self._impl).revision, raise_sync_error=False)(message, file_type)
         else:
-            cast("SyncMigrationCommands[Any]", self._impl).revision(message)
+            cast("SyncMigrationCommands[Any]", self._impl).revision(message, file_type)