Save user version in step-by-step migrations

drift/CHANGELOG.md

@@ -3,6 +3,7 @@
 - Add APIs to setup Wasm databases with custom drift workers.
 - Add `Expression.and` and `Expression.or` to create disjunctions and conjunctions
   of sub-predicates.
+- Step-by-step migrations now save the intermediate schema version after each step.
 
 ## 2.12.1
 

drift/lib/internal/versioned_schema.dart

@@ -103,6 +103,16 @@ abstract base class VersionedSchema {
       final newVersion = await steps(target, database);
       assert(newVersion > target);
 
+      // Saving the schema version after each step prevents the schema of the
+      // database diverging from what's stored in `user_version` if a migration
+      // fails halfway.
+      // We can only reliably do this for sqlite3 at the moment since managing
+      // schema versions happens at a lower layer and is not current exposed to
+      // the query builder.
+      if (database.executor.dialect == SqlDialect.sqlite) {
+        await database.customStatement('pragma user_version = $newVersion');
+      }
+
       target = newVersion;
     }
   }

drift/test/integration_tests/migrations_integration_test.dart

@@ -1,5 +1,6 @@
 @TestOn('vm')
 import 'package:drift/drift.dart' hide isNull;
+import 'package:drift/internal/versioned_schema.dart';
 import 'package:drift/native.dart';
 import 'package:drift_dev/api/migrations.dart';
 import 'package:sqlite3/sqlite3.dart';
@@ -461,6 +462,36 @@ void main() {
       }
     }
   });
+
+  test('step-by-step migrations saves state halfway', () async {
+    final underlying = sqlite3.openInMemory()
+      ..execute('pragma user_version = 1;');
+    addTearDown(underlying.dispose);
+
+    final expectedException = Exception('schema upgrade!');
+
+    final db = TodoDb(NativeDatabase.opened(underlying))
+      ..schemaVersion = 5
+      ..migration =
+          MigrationStrategy(onUpgrade: VersionedSchema.stepByStepHelper(
+        step: (version, database) async {
+          await database.customStatement('CREATE TABLE t$version (id INT);');
+
+          if (version == 3) {
+            throw expectedException;
+          }
+
+          return version + 1;
+        },
+      ));
+
+    await expectLater(
+      db.customSelect('SELECT 1').get(),
+      throwsA(expectedException),
+    );
+
+    expect(underlying.userVersion, 3);
+  });
 }
 
 class _TestDatabase extends GeneratedDatabase {