Add API for finegrained control for steps by steps

docs/lib/snippets/migrations/schema_versions.dart

@@ -106,11 +106,11 @@ class Shape1 extends i0.VersionedTable {
 i1.GeneratedColumn<int> _column_5(String aliasedName) =>
     i1.GeneratedColumn<int>('priority', aliasedName, true,
         type: i1.DriftSqlType.int);
-i1.OnUpgrade stepByStep({
+i0.MigrationStepWithVersion migrationSteps({
   required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
   required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
 }) {
-  return i1.Migrator.stepByStepHelper(step: (currentVersion, database) async {
+  return (currentVersion, database) async {
     switch (currentVersion) {
       case 1:
         final schema = _S2(database: database);
@@ -125,5 +125,15 @@ i1.OnUpgrade stepByStep({
       default:
         throw ArgumentError.value('Unknown migration from $currentVersion');
     }
-  });
+  };
 }
+
+i1.OnUpgrade stepByStep({
+  required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
+  required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
+}) =>
+    i0.VersionedSchema.stepByStepHelper(
+        step: migrationSteps(
+      from1To2: from1To2,
+      from2To3: from2To3,
+    ));

drift/lib/internal/versioned_schema.dart

@@ -12,6 +12,15 @@ library;
 
 import 'package:drift/drift.dart';
 
+/// Signature of a function, typically generated by drift, that runs a single
+/// migration step with a given [currentVersion] and the [database].
+///
+/// Returns the schema version code that the function migrates to.
+typedef MigrationStepWithVersion = Future<int> Function(
+  int currentVersion,
+  GeneratedDatabase database,
+);
+
 /// A snapshot of a database schema at a previous version.
 ///
 /// This class is meant to be extended by generated code.
@@ -27,6 +36,77 @@ abstract base class VersionedSchema {
 
   /// All drift schema entities at the time of the set [version].
   Iterable<DatabaseSchemaEntity> get entities;
+
+  /// A helper used by drift internally to implement the [step-by-step](https://drift.simonbinder.eu/docs/advanced-features/migrations/#step-by-step)
+  /// migration feature.
+  ///
+  /// This method implements an [OnUpgrade] callback by repeatedly invoking
+  /// [step] with the current version, assuming that [step] will perform an
+  /// upgrade from that version to the version returned by the callback.
+  ///
+  /// If you want to customize the way the migration steps are invoked, for
+  /// instance by running statements before and afterwards, see
+  /// [runMigrationSteps].
+  static OnUpgrade stepByStepHelper({
+    required MigrationStepWithVersion step,
+  }) {
+    return (m, from, to) async {
+      return await runMigrationSteps(
+        migrator: m,
+        from: from,
+        to: to,
+        steps: step,
+      );
+    };
+  }
+
+  /// Helper method that runs a (subset of) [stepByStepHelper] by invoking the
+  /// [steps] function for each intermediate schema version from [from] until
+  /// [to] is reached.
+  ///
+  /// This can be used to implement a custom `OnUpgrade` callback that runs
+  /// additional checks before and after the migrations:
+  ///
+  /// ```dart
+  /// onUpgrade: (m, from, to) async {
+  ///  await customStatement('PRAGMA foreign_keys = OFF');
+  ///
+  ///  await transaction(
+  ///    () => VersionedSchema.runMigrationSteps(
+  ///      migrator: m,
+  ///      from: from,
+  ///      to: to,
+  ///      steps: migrationSteps(
+  ///        from1To2: ...,
+  ///        ...
+  ///      ),
+  ///    ),
+  ///  );
+  ///
+  ///  if (kDebugMode) {
+  ///    final wrongForeignKeys = await customSelect('PRAGMA foreign_key_check').get();
+  ///    assert(wrongForeignKeys.isEmpty, '${wrongForeignKeys.map((e) => e.data)}');
+  ///  }
+  ///
+  ///  await customStatement('PRAGMA foreign_keys = ON;');
+  /// },
+  /// ```
+
+  static Future<void> runMigrationSteps({
+    required Migrator migrator,
+    required int from,
+    required int to,
+    required MigrationStepWithVersion steps,
+  }) async {
+    final database = migrator.database;
+
+    for (var target = from; target < to;) {
+      final newVersion = await steps(target, database);
+      assert(newVersion > target);
+
+      target = newVersion;
+    }
+  }
 }
 
 /// A drift table implementation that, instead of being generated, is constructed

drift/lib/src/runtime/query_builder/migration.dart

@@ -49,17 +49,20 @@ class MigrationStrategy {
 
 /// Runs migrations declared by a [MigrationStrategy].
 class Migrator {
-  final GeneratedDatabase _db;
+  /// The instance of the [GeneratedDatabase] class generated by `drift_dev`
+  /// that this migrator is connected to.
+  final GeneratedDatabase database;
+
+  /// When non-null, use an old schema version in [createAll] and similar
+  /// methods, making it easier to write sound migrations that don't always
+  /// assume the latest database schema.
   final VersionedSchema? _fixedVersion;
 
   /// Used internally by drift when opening the database.
-  Migrator(this._db, [this._fixedVersion]);
+  Migrator(this.database, [this._fixedVersion]);
 
   Iterable<DatabaseSchemaEntity> get _allSchemaEntities {
-    return switch (_fixedVersion) {
+    return _fixedVersion?.entities ?? database.allSchemaEntities;
-      null => _db.allSchemaEntities,
-      var fixed => fixed.entities,
-    };
   }
 
   /// Creates all tables specified for the database, if they don't exist
@@ -112,7 +115,8 @@ class Migrator {
   }
 
   GenerationContext _createContext({bool supportsVariables = false}) {
-    return GenerationContext.fromDb(_db, supportsVariables: supportsVariables);
+    return GenerationContext.fromDb(database,
+        supportsVariables: supportsVariables);
   }
 
   /// Creates the given table if it doesn't exist
@@ -152,26 +156,26 @@ class Migrator {
   /// [drift docs]: https://drift.simonbinder.eu/docs/advanced-features/migrations/#complex-migrations
   Future<void> alterTable(TableMigration migration) async {
     final foreignKeysEnabled =
-        (await _db.customSelect('PRAGMA foreign_keys').getSingle())
+        (await database.customSelect('PRAGMA foreign_keys').getSingle())
             .read<bool>('foreign_keys');
     final legacyAlterTable =
-        (await _db.customSelect('PRAGMA legacy_alter_table').getSingle())
+        (await database.customSelect('PRAGMA legacy_alter_table').getSingle())
             .read<bool>('legacy_alter_table');
 
     if (foreignKeysEnabled) {
-      await _db.customStatement('PRAGMA foreign_keys = OFF;');
+      await database.customStatement('PRAGMA foreign_keys = OFF;');
     }
 
     final table = migration.affectedTable;
     final tableName = table.actualTableName;
 
-    await _db.transaction(() async {
+    await database.transaction(() async {
       // We will drop the original table later, which will also delete
       // associated triggers, indices and and views. We query sqlite_schema to
       // re-create those later.
       // We use the legacy sqlite_master table since the _schema rename happened
       // in a very recent version (3.33.0)
-      final schemaQuery = await _db.customSelect(
+      final schemaQuery = await database.customSelect(
         'SELECT type, name, sql FROM sqlite_master WHERE tbl_name = ?;',
         variables: [Variable<String>(tableName)],
       ).get();
@@ -274,7 +278,7 @@ class Migrator {
 
     // Finally, re-enable foreign keys if they were enabled originally.
     if (foreignKeysEnabled) {
-      await _db.customStatement('PRAGMA foreign_keys = ON;');
+      await database.customStatement('PRAGMA foreign_keys = ON;');
     }
   }
 
@@ -378,7 +382,8 @@ class Migrator {
     if (stmts != null) {
       await _issueQueryByDialect(stmts);
     } else if (view.query != null) {
-      final context = GenerationContext.fromDb(_db, supportsVariables: false);
+      final context =
+          GenerationContext.fromDb(database, supportsVariables: false);
       final columnNames = view.$columns
           .map((e) => e.escapedNameFor(context.dialect))
           .join(', ');
@@ -493,7 +498,7 @@ class Migrator {
   }
 
   Future<void> _issueCustomQuery(String sql, [List<dynamic>? args]) {
-    return _db.customStatement(sql, args);
+    return database.customStatement(sql, args);
   }
 
   /// A helper used by drift internally to implement the [step-by-step](https://drift.simonbinder.eu/docs/advanced-features/migrations/#step-by-step)
@@ -502,23 +507,12 @@ class Migrator {
   /// This method implements an [OnUpgrade] callback by repeatedly invoking
   /// [step] with the current version, assuming that [step] will perform an
   /// upgrade from that version to the version returned by the callback.
-  @experimental
+  @Deprecated(
+      'Re-generate code so that it uses `VersionedSchema.stepByStepHelper`')
   static OnUpgrade stepByStepHelper({
-    required Future<int> Function(
+    required MigrationStepWithVersion step,
-      int currentVersion,
-      GeneratedDatabase database,
-    ) step,
   }) {
-    return (m, from, to) async {
+    return VersionedSchema.stepByStepHelper(step: step);
-      final database = m._db;
-
-      for (var target = from; target < to;) {
-        final newVersion = await step(target, database);
-        assert(newVersion > target);
-
-        target = newVersion;
-      }
-    };
   }
 }
 
@@ -564,7 +558,7 @@ extension DestructiveMigrationExtension on GeneratedDatabase {
       onUpgrade: (m, from, to) async {
         // allSchemaEntities are sorted topologically references between them.
         // Reverse order for deletion in order to not break anything.
-        final reversedEntities = m._db.allSchemaEntities.toList().reversed;
+        final reversedEntities = m._allSchemaEntities.toList().reversed;
 
         for (final entity in reversedEntities) {
           await m.drop(entity);

drift_dev/lib/src/writer/schema_version_writer.dart

@@ -297,6 +297,16 @@ class SchemaVersionWriter {
     return name;
   }
 
+  void _writeCallbackArgsForStep(TextEmitter text) {
+    for (final (current, next) in versions.withNext) {
+      text
+        ..write('required Future<void> Function(')
+        ..writeDriftRef('Migrator')
+        ..write(' m, _S${next.version} schema)')
+        ..writeln('from${current.version}To${next.version},');
+    }
+  }
+
   void write() {
     libraryScope.leaf()
       ..writeln('// ignore_for_file: type=lint,unused_import')
@@ -337,31 +347,21 @@ class SchemaVersionWriter {
       versionScope.leaf().writeln('}');
     }
 
-    // Write a stepByStep migration function that takes a callback doing a step
+    // Write a MigrationStepWithVersion factory that takes a callback doing a
-    // for each schema to the next. We supply a special migrator that only
+    // step for each schema to to the next. We supply a special migrator that
-    // considers entities from that version, as well as a typed reference to the
+    // only considers entities from that version, as well as a typed reference
-    // _S<x> class used to lookup elements.
+    // to the _S<x> class used to lookup elements.
-    final stepByStep = libraryScope.leaf()
+    final steps = libraryScope.leaf()
-      ..writeDriftRef('OnUpgrade')
+      ..writeUriRef(_schemaLibrary, 'MigrationStepWithVersion')
-      ..write(' stepByStep({');
+      ..write(' migrationSteps({');
-
+    _writeCallbackArgsForStep(steps);
-    for (final (current, next) in versions.withNext) {
+    steps
-      stepByStep
-        ..write('required Future<void> Function(')
-        ..writeDriftRef('Migrator')
-        ..write(' m, _S${next.version} schema)')
-        ..writeln('from${current.version}To${next.version},');
-    }
-
-    stepByStep
       ..writeln('}) {')
-      ..write('return ')
+      ..writeln('return (currentVersion, database) async {')
-      ..writeDriftRef('Migrator')
-      ..writeln('.stepByStepHelper(step: (currentVersion, database) async {')
       ..writeln('switch (currentVersion) {');
 
     for (final (current, next) in versions.withNext) {
-      stepByStep
+      steps
         ..writeln('case ${current.version}:')
         ..write('final schema = _S${next.version}(database: database);')
         ..write('final migrator = ')
@@ -372,13 +372,29 @@ class SchemaVersionWriter {
         ..writeln('return ${next.version};');
     }
 
-    stepByStep
+    steps
       ..writeln(
           r"default: throw ArgumentError.value('Unknown migration from $currentVersion');")
       ..writeln('}') // End of switch
-      ..writeln('}') // End of stepByStepHelper function
+      ..writeln('};') // End of function literal
-      ..writeln(');') // End of stepByStepHelper call
+      ..writeln('}'); // End of migrationSteps method
-      ..writeln('}'); // End of method
+
+    final stepByStep = libraryScope.leaf()
+      ..writeDriftRef('OnUpgrade')
+      ..write(' stepByStep({');
+    _writeCallbackArgsForStep(stepByStep);
+    stepByStep
+      ..writeln('}) => ')
+      ..writeUriRef(_schemaLibrary, 'VersionedSchema')
+      ..write('.stepByStepHelper(step: migrationSteps(');
+
+    for (final (current, next) in versions.withNext) {
+      final name = 'from${current.version}To${next.version}';
+
+      stepByStep.writeln('$name: $name,');
+    }
+
+    stepByStep.writeln('));');
   }
 }
 

examples/app/lib/database/schema_versions.dart

@@ -179,11 +179,11 @@ i1.GeneratedColumn<int> _column_7(String aliasedName) =>
         type: i1.DriftSqlType.int,
         defaultConstraints: i1.GeneratedColumn.constraintIsAlways(
             'REFERENCES categories (id)'));
-i1.OnUpgrade stepByStep({
+i0.MigrationStepWithVersion migrationSteps({
   required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
   required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
 }) {
-  return i1.Migrator.stepByStepHelper(step: (currentVersion, database) async {
+  return (currentVersion, database) async {
     switch (currentVersion) {
       case 1:
         final schema = _S2(database: database);
@@ -198,5 +198,15 @@ i1.OnUpgrade stepByStep({
       default:
         throw ArgumentError.value('Unknown migration from $currentVersion');
     }
-  });
+  };
 }
+
+i1.OnUpgrade stepByStep({
+  required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
+  required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
+}) =>
+    i0.VersionedSchema.stepByStepHelper(
+        step: migrationSteps(
+      from1To2: from1To2,
+      from2To3: from2To3,
+    ));

examples/migrations_example/lib/database.dart

@@ -1,4 +1,5 @@
 import 'package:drift/drift.dart';
+import 'package:drift/internal/versioned_schema.dart';
 import 'package:drift_dev/api/migrations.dart';
 
 import 'tables.dart';
@@ -6,6 +7,10 @@ import 'src/versions.dart';
 
 part 'database.g.dart';
 
+/// This isn't a Flutter app, so we have to define the constant. In a real app,
+/// just use the constant defined in the Flutter SDK.
+const kDebugMode = true;
+
 @DriftDatabase(include: {'tables.drift'})
 class Database extends _$Database {
   static const latestSchemaVersion = 9;
@@ -18,7 +23,28 @@ class Database extends _$Database {
   @override
   MigrationStrategy get migration {
     return MigrationStrategy(
-      onUpgrade: _upgrade,
+      onUpgrade: (m, from, to) async {
+        // Following the advice from https://drift.simonbinder.eu/docs/advanced-features/migrations/#tips
+        await customStatement('PRAGMA foreign_keys = OFF');
+
+        await transaction(
+          () => VersionedSchema.runMigrationSteps(
+            migrator: m,
+            from: from,
+            to: to,
+            steps: _upgrade,
+          ),
+        );
+
+        if (kDebugMode) {
+          final wrongForeignKeys =
+              await customSelect('PRAGMA foreign_key_check').get();
+          assert(wrongForeignKeys.isEmpty,
+              '${wrongForeignKeys.map((e) => e.data)}');
+        }
+
+        await customStatement('PRAGMA foreign_keys = ON');
+      },
       beforeOpen: (details) async {
         // For Flutter apps, this should be wrapped in an if (kDebugMode) as
         // suggested here: https://drift.simonbinder.eu/docs/advanced-features/migrations/#verifying-a-database-schema-at-runtime
@@ -27,7 +53,7 @@ class Database extends _$Database {
     );
   }
 
-  static final _upgrade = stepByStep(
+  static final _upgrade = migrationSteps(
     from1To2: (m, schema) async {
       // Migration from 1 to 2: Add name column in users. Use "no name"
       // as a default value.

examples/migrations_example/lib/src/versions.dart

@@ -585,7 +585,7 @@ i1.GeneratedColumn<int> _column_17(String aliasedName) =>
     i1.GeneratedColumn<int>('owner', aliasedName, false,
         type: i1.DriftSqlType.int,
         $customConstraints: 'NOT NULL REFERENCES users(id)');
-i1.OnUpgrade stepByStep({
+i0.MigrationStepWithVersion migrationSteps({
   required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
   required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
   required Future<void> Function(i1.Migrator m, _S4 schema) from3To4,
@@ -595,7 +595,7 @@ i1.OnUpgrade stepByStep({
   required Future<void> Function(i1.Migrator m, _S8 schema) from7To8,
   required Future<void> Function(i1.Migrator m, _S9 schema) from8To9,
 }) {
-  return i1.Migrator.stepByStepHelper(step: (currentVersion, database) async {
+  return (currentVersion, database) async {
     switch (currentVersion) {
       case 1:
         final schema = _S2(database: database);
@@ -640,5 +640,27 @@ i1.OnUpgrade stepByStep({
       default:
         throw ArgumentError.value('Unknown migration from $currentVersion');
     }
-  });
+  };
 }
+
+i1.OnUpgrade stepByStep({
+  required Future<void> Function(i1.Migrator m, _S2 schema) from1To2,
+  required Future<void> Function(i1.Migrator m, _S3 schema) from2To3,
+  required Future<void> Function(i1.Migrator m, _S4 schema) from3To4,
+  required Future<void> Function(i1.Migrator m, _S5 schema) from4To5,
+  required Future<void> Function(i1.Migrator m, _S6 schema) from5To6,
+  required Future<void> Function(i1.Migrator m, _S7 schema) from6To7,
+  required Future<void> Function(i1.Migrator m, _S8 schema) from7To8,
+  required Future<void> Function(i1.Migrator m, _S9 schema) from8To9,
+}) =>
+    i0.VersionedSchema.stepByStepHelper(
+        step: migrationSteps(
+      from1To2: from1To2,
+      from2To3: from2To3,
+      from3To4: from3To4,
+      from4To5: from4To5,
+      from5To6: from5To6,
+      from6To7: from6To7,
+      from7To8: from7To8,
+      from8To9: from8To9,
+    ));