feat: add transactional migrations support (#132) (#137)

cofin · web-flow · commit 67ca8b5ff2a2 · 2025-10-19T17:52:03.000-05:00
Adds transactional migrations for database adapters that support DDL transactions. When enabled, migrations run atomically - either everything succeeds or everything rolls back. No more partial migrations leaving your database in an inconsistent state. Builds on the hybrid versioning work in #128.
diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py
@@ -100,6 +100,7 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]):
 
     driver_type: ClassVar[type[AdbcDriver]] = AdbcDriver
     connection_type: "ClassVar[type[AdbcConnection]]" = AdbcConnection
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/aiosqlite/config.py b/sqlspec/adapters/aiosqlite/config.py
@@ -73,6 +73,7 @@ class AiosqliteConfig(AsyncDatabaseConfig["AiosqliteConnection", AiosqliteConnec
 
     driver_type: "ClassVar[type[AiosqliteDriver]]" = AiosqliteDriver
     connection_type: "ClassVar[type[AiosqliteConnection]]" = AiosqliteConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py
@@ -83,6 +83,7 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
 
     driver_type: ClassVar[type[AsyncmyDriver]] = AsyncmyDriver
     connection_type: "ClassVar[type[AsyncmyConnection]]" = AsyncmyConnection  # pyright: ignore
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py
@@ -96,6 +96,7 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
 
     driver_type: "ClassVar[type[AsyncpgDriver]]" = AsyncpgDriver
     connection_type: "ClassVar[type[AsyncpgConnection]]" = type(AsyncpgConnection)  # type: ignore[assignment]
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py
@@ -88,6 +88,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
 
     driver_type: ClassVar[type[BigQueryDriver]] = BigQueryDriver
     connection_type: "ClassVar[type[BigQueryConnection]]" = BigQueryConnection
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/duckdb/config.py b/sqlspec/adapters/duckdb/config.py
@@ -176,6 +176,7 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
 
     driver_type: "ClassVar[type[DuckDBDriver]]" = DuckDBDriver
     connection_type: "ClassVar[type[DuckDBConnection]]" = DuckDBConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py
@@ -106,6 +106,7 @@ class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "OracleSyncConne
     driver_type: ClassVar[type[OracleSyncDriver]] = OracleSyncDriver
     connection_type: "ClassVar[type[OracleSyncConnection]]" = OracleSyncConnection
     migration_tracker_type: "ClassVar[type[OracleSyncMigrationTracker]]" = OracleSyncMigrationTracker
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -263,6 +264,7 @@ class OracleAsyncConfig(AsyncDatabaseConfig[OracleAsyncConnection, "OracleAsyncC
     connection_type: "ClassVar[type[OracleAsyncConnection]]" = OracleAsyncConnection
     driver_type: ClassVar[type[OracleAsyncDriver]] = OracleAsyncDriver
     migration_tracker_type: "ClassVar[type[OracleAsyncMigrationTracker]]" = OracleAsyncMigrationTracker
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py
@@ -94,6 +94,7 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
 
     driver_type: ClassVar[type[PsqlpyDriver]] = PsqlpyDriver
     connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py
@@ -95,6 +95,7 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
 
     driver_type: "ClassVar[type[PsycopgSyncDriver]]" = PsycopgSyncDriver
     connection_type: "ClassVar[type[PsycopgSyncConnection]]" = PsycopgSyncConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
@@ -285,6 +286,7 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
 
     driver_type: ClassVar[type[PsycopgAsyncDriver]] = PsycopgAsyncDriver
     connection_type: "ClassVar[type[PsycopgAsyncConnection]]" = PsycopgAsyncConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/adapters/sqlite/config.py b/sqlspec/adapters/sqlite/config.py
@@ -62,6 +62,7 @@ class SqliteConfig(SyncDatabaseConfig[SqliteConnection, SqliteConnectionPool, Sq
 
     driver_type: "ClassVar[type[SqliteDriver]]" = SqliteDriver
     connection_type: "ClassVar[type[SqliteConnection]]" = SqliteConnection
+    supports_transactional_ddl: "ClassVar[bool]" = True
 
     def __init__(
         self,
diff --git a/sqlspec/config.py b/sqlspec/config.py
@@ -94,6 +94,9 @@ class MigrationConfig(TypedDict):
     include_extensions: NotRequired["list[str]"]
     """List of extension names whose migrations should be included. Extension migrations maintain separate versioning and are prefixed with 'ext_{name}_'."""
 
+    transactional: NotRequired[bool]
+    """Wrap migrations in transactions when supported. When enabled (default for adapters that support it), each migration runs in a transaction that is committed on success or rolled back on failure. This prevents partial migrations from leaving the database in an inconsistent state. Requires adapter support for transactional DDL. Defaults to True for PostgreSQL, SQLite, and DuckDB; False for MySQL, Oracle, and BigQuery. Individual migrations can override this with a '-- transactional: false' comment."""
+
 
 class LitestarConfig(TypedDict):
     """Configuration options for Litestar SQLSpec plugin.
@@ -260,6 +263,7 @@ class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
     connection_type: "ClassVar[type[Any]]"
     is_async: "ClassVar[bool]" = False
     supports_connection_pooling: "ClassVar[bool]" = False
+    supports_transactional_ddl: "ClassVar[bool]" = False
     supports_native_arrow_import: "ClassVar[bool]" = False
     supports_native_arrow_export: "ClassVar[bool]" = False
     supports_native_parquet_import: "ClassVar[bool]" = False
diff --git a/sqlspec/migrations/commands.py b/sqlspec/migrations/commands.py
@@ -268,14 +268,19 @@ def upgrade(
                     continue
 
                 try:
-                    _, execution_time = self.runner.execute_upgrade(driver, migration)
-                    self.tracker.record_migration(
-                        driver, migration["version"], migration["description"], execution_time, migration["checksum"]
-                    )
+
+                    def record_version(exec_time: int, migration: "dict[str, Any]" = migration) -> None:
+                        self.tracker.record_migration(
+                            driver, migration["version"], migration["description"], exec_time, migration["checksum"]
+                        )
+
+                    _, execution_time = self.runner.execute_upgrade(driver, migration, on_success=record_version)
                     console.print(f"[green]✓ Applied in {execution_time}ms[/]")
 
                 except Exception as e:
-                    console.print(f"[red]✗ Failed: {e}[/]")
+                    use_txn = self.runner.should_use_transaction(migration, self.config)
+                    rollback_msg = " (transaction rolled back)" if use_txn else ""
+                    console.print(f"[red]✗ Failed{rollback_msg}: {e}[/]")
                     return
 
             if dry_run:
@@ -332,11 +337,16 @@ def downgrade(self, revision: str = "-1", *, dry_run: bool = False) -> None:
                     continue
 
                 try:
-                    _, execution_time = self.runner.execute_downgrade(driver, migration)
-                    self.tracker.remove_migration(driver, version)
+
+                    def remove_version(exec_time: int, version: str = version) -> None:
+                        self.tracker.remove_migration(driver, version)
+
+                    _, execution_time = self.runner.execute_downgrade(driver, migration, on_success=remove_version)
                     console.print(f"[green]✓ Reverted in {execution_time}ms[/]")
                 except Exception as e:
-                    console.print(f"[red]✗ Failed: {e}[/]")
+                    use_txn = self.runner.should_use_transaction(migration, self.config)
+                    rollback_msg = " (transaction rolled back)" if use_txn else ""
+                    console.print(f"[red]✗ Failed{rollback_msg}: {e}[/]")
                     return
 
             if dry_run:
@@ -697,13 +707,18 @@ async def upgrade(
                     continue
 
                 try:
-                    _, execution_time = await self.runner.execute_upgrade(driver, migration)
-                    await self.tracker.record_migration(
-                        driver, migration["version"], migration["description"], execution_time, migration["checksum"]
-                    )
+
+                    async def record_version(exec_time: int, migration: "dict[str, Any]" = migration) -> None:
+                        await self.tracker.record_migration(
+                            driver, migration["version"], migration["description"], exec_time, migration["checksum"]
+                        )
+
+                    _, execution_time = await self.runner.execute_upgrade(driver, migration, on_success=record_version)
                     console.print(f"[green]✓ Applied in {execution_time}ms[/]")
                 except Exception as e:
-                    console.print(f"[red]✗ Failed: {e}[/]")
+                    use_txn = self.runner.should_use_transaction(migration, self.config)
+                    rollback_msg = " (transaction rolled back)" if use_txn else ""
+                    console.print(f"[red]✗ Failed{rollback_msg}: {e}[/]")
                     return
 
             if dry_run:
@@ -761,11 +776,18 @@ async def downgrade(self, revision: str = "-1", *, dry_run: bool = False) -> Non
                     continue
 
                 try:
-                    _, execution_time = await self.runner.execute_downgrade(driver, migration)
-                    await self.tracker.remove_migration(driver, version)
+
+                    async def remove_version(exec_time: int, version: str = version) -> None:
+                        await self.tracker.remove_migration(driver, version)
+
+                    _, execution_time = await self.runner.execute_downgrade(
+                        driver, migration, on_success=remove_version
+                    )
                     console.print(f"[green]✓ Reverted in {execution_time}ms[/]")
                 except Exception as e:
-                    console.print(f"[red]✗ Failed: {e}[/]")
+                    use_txn = self.runner.should_use_transaction(migration, self.config)
+                    rollback_msg = " (transaction rolled back)" if use_txn else ""
+                    console.print(f"[red]✗ Failed{rollback_msg}: {e}[/]")
                     return
 
             if dry_run:
diff --git a/sqlspec/migrations/context.py b/sqlspec/migrations/context.py
@@ -129,9 +129,6 @@ def validate_async_usage(self, migration_func: Any) -> None:
 
         Args:
             migration_func: The migration function to validate.
-
-        Raises:
-            RuntimeError: If async function is used inappropriately.
         """
         if inspect.iscoroutinefunction(migration_func) and not self.is_async_execution and not self.is_async_driver:
             msg = (
diff --git a/sqlspec/migrations/runner.py b/sqlspec/migrations/runner.py
