Run each migration inside a transaction

Migrations.apply() now wraps each migration function and its
_sqlite_migrations tracking row in db.atomic(), so they commit
together: a failing migration rolls back cleanly, is not recorded,
and stays pending, eliminating the double-apply hazard where a
partially-failed migration left committed side effects behind.

Migrations that cannot run inside a transaction (VACUUM, journal
mode changes, manual transaction management) can opt out by
registering with @migrations(transactional=False).

Also fixed the test fixtures which used db.query() for INSERT
statements - query() is a lazy generator, so those inserts never
actually executed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UnLnhsH25Nnv7LHhekUfPd

docs/migrations.rst

@@ -70,6 +70,25 @@ When you apply a set of migrations you can stop part way through by specifying a
 
     migrations.apply(db, stop_before="add_weight")
 
+.. _migrations_transactions:
+
+Migrations and transactions
+===========================
+
+Each migration runs inside a transaction, together with the ``_sqlite_migrations`` record of it having been applied. If a migration function raises an exception, everything it did is rolled back, no record is written and the migration stays pending - so fixing the error and re-applying will run that migration again from a clean state. Migrations that completed earlier in the same ``apply()`` run stay applied.
+
+Some operations cannot run inside a transaction, for example ``VACUUM`` or changing the journal mode with ``db.enable_wal()``. Register migrations like these with ``transactional=False``:
+
+.. code-block:: python
+
+    @migrations(transactional=False)
+    def compact(db):
+        db.execute("VACUUM")
+
+A migration registered with ``transactional=False`` runs without a wrapping transaction, so if it fails part way through any changes it already made will not be rolled back, and re-applying will run the whole function again.
+
+Avoid calling ``db.conn.commit()`` or otherwise managing transactions manually inside a transactional migration - register the migration with ``transactional=False`` if it needs to control its own transactions.
+
 Applying migrations using the CLI
 =================================
 

sqlite_utils/migrations.py

@@ -14,6 +14,7 @@ class Migrations:
     class _Migration:
         name: str
         fn: Callable
+        transactional: bool = True
 
     @dataclass
     class _AppliedMigration:
@@ -27,15 +28,22 @@ class Migrations:
         self.name = name
         self._migrations: list[Migrations._Migration] = []
 
-    def __call__(self, *, name: str | None = None) -> Callable:
+    def __call__(
+        self, *, name: str | None = None, transactional: bool = True
+    ) -> Callable:
         """
         :param name: The name to use for this migration - if not provided,
           the name of the function will be used.
+        :param transactional: If ``True`` (the default) the migration and the
+          record of it having been applied are wrapped in a transaction, which
+          will be rolled back if the migration raises an exception. Pass
+          ``False`` for migrations that cannot run inside a transaction, for
+          example those that execute ``VACUUM``.
         """
 
         def inner(func: Callable) -> Callable:
             self._migrations.append(
-                self._Migration(name or getattr(func, "__name__"), func)
+                self._Migration(name or getattr(func, "__name__"), func, transactional)
             )
             return func
 
@@ -73,6 +81,12 @@ class Migrations:
     def apply(self, db: "Database", *, stop_before: str | Iterable[str] | None = None):
         """
         Apply any pending migrations to the database.
+
+        Each migration runs inside a transaction, together with the record of
+        it having been applied - if the migration raises an exception its
+        changes are rolled back, no record is written and the migration stays
+        pending. Migrations registered with ``transactional=False`` run
+        outside of a transaction.
         """
         self.ensure_migrations_table(db)
         if stop_before is None:
@@ -85,14 +99,22 @@ class Migrations:
             name = migration.name
             if name in stop_before_names:
                 return
-            migration.fn(db)
-            _table(db, self.migrations_table).insert(
-                {
-                    "migration_set": self.name,
-                    "name": name,
-                    "applied_at": str(datetime.datetime.now(datetime.timezone.utc)),
-                }
-            )
+            if migration.transactional:
+                with db.atomic():
+                    migration.fn(db)
+                    self._record_applied(db, name)
+            else:
+                migration.fn(db)
+                self._record_applied(db, name)
+
+    def _record_applied(self, db: "Database", name: str):
+        _table(db, self.migrations_table).insert(
+            {
+                "migration_set": self.name,
+                "name": name,
+                "applied_at": str(datetime.datetime.now(datetime.timezone.utc)),
+            }
+        )
 
     def ensure_migrations_table(self, db: "Database"):
         """

tests/test_migrations.py

@@ -14,7 +14,7 @@ def migrations():
     @migrations()
     def m002(db):
         db["cats"].create({"name": str})
-        db.query("insert into dogs (name) values ('Pancakes')")
+        db.execute("insert into dogs (name) values ('Pancakes')")
 
     return migrations
 
@@ -32,7 +32,7 @@ def migrations_not_ordered_alphabetically():
     @migrations()
     def m001(db):
         db["cats"].create({"name": str})
-        db.query("insert into dogs (name) values ('Pancakes')")
+        db.execute("insert into dogs (name) values ('Pancakes')")
 
     return migrations
 
@@ -80,6 +80,76 @@ def test_order_does_not_matter(migrations, migrations_not_ordered_alphabetically
     assert db1.schema == db2.schema
 
 
+def test_failing_migration_rolls_back(migrations):
+    @migrations()
+    def m003(db):
+        db["birds"].create({"name": str})
+        db.execute("insert into dogs (name) values ('Dozer')")
+        raise ValueError("boom")
+
+    db = sqlite_utils.Database(memory=True)
+    with pytest.raises(ValueError):
+        migrations.apply(db)
+    # m001 and m002 committed before the failure and stay applied
+    assert set(db.table_names()) == {"_sqlite_migrations", "dogs", "cats"}
+    assert [r["name"] for r in db["dogs"].rows] == ["Cleo", "Pancakes"]
+    assert [m.name for m in migrations.applied(db)] == ["m001", "m002"]
+    # Everything m003 did was rolled back and it is still pending
+    assert [m.name for m in migrations.pending(db)] == ["m003"]
+
+
+def test_rerun_after_failure_applies_each_migration_once():
+    state = {"fail": True}
+    migrations = Migrations("test")
+
+    @migrations()
+    def m001(db):
+        db["dogs"].insert({"name": "Cleo"})
+
+    @migrations()
+    def m002(db):
+        db["dogs"].insert({"name": "Pancakes"})
+        if state["fail"]:
+            raise ValueError("boom")
+
+    db = sqlite_utils.Database(memory=True)
+    with pytest.raises(ValueError):
+        migrations.apply(db)
+    state["fail"] = False
+    migrations.apply(db)
+    # m001 must not have been re-applied, m002 applied exactly once
+    assert [r["name"] for r in db["dogs"].rows] == ["Cleo", "Pancakes"]
+    assert [m.name for m in migrations.applied(db)] == ["m001", "m002"]
+
+
+def test_non_transactional_migration_allows_vacuum(tmpdir):
+    path = str(tmpdir / "test.db")
+    db = sqlite_utils.Database(path)
+    migrations = Migrations("test")
+
+    @migrations()
+    def m001(db):
+        db["dogs"].insert({"name": "Cleo"})
+
+    @migrations(transactional=False)
+    def m002(db):
+        db.execute("VACUUM")
+
+    migrations.apply(db)
+    assert [m.name for m in migrations.applied(db)] == ["m001", "m002"]
+    db.close()
+
+
+def test_apply_composes_inside_outer_transaction(migrations):
+    db = sqlite_utils.Database(memory=True)
+    with pytest.raises(ZeroDivisionError):
+        with db.atomic():
+            migrations.apply(db)
+            raise ZeroDivisionError
+    # The outer transaction rolled back, taking the migrations with it
+    assert db.table_names() == []
+
+
 @pytest.mark.parametrize(
     "create_table,pk",
     (