Allow registering DriverInitializer for VerifyMigrationTask with ServiceLoader mechanism

sqldelight-gradle-plugin/src/main/kotlin/app/cash/sqldelight/gradle/VerifyMigrationTask.kt

@@ -29,6 +29,8 @@ import org.gradle.api.tasks.TaskAction
 import org.gradle.workers.WorkAction
 import org.gradle.workers.WorkParameters
 import java.io.File
+import java.io.Serializable
+import java.sql.DriverManager
 import java.util.ServiceLoader
 
 @CacheableTask
@@ -51,6 +53,8 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
 
   @Input var verifyDefinitions: Boolean = true
 
+  @get:Internal var driverInitializer: DriverInitializer? = null
+
   /* Tasks without an output are never considered UP-TO-DATE by Gradle. Adding an output file that's created when the
    * task completes successfully works around the lack of an output for this task. There may be a better solution once
    * https://github.com/gradle/gradle/issues/14223 is resolved. */
@@ -68,6 +72,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
         it.verifyMigrations.set(verifyMigrations)
         it.compilationUnit.set(compilationUnit)
         it.verifyDefinitions.set(verifyDefinitions)
+        it.driverInitializer.set(driverInitializer)
       }
       workQueue.await()
     }.onSuccess {
@@ -92,6 +97,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
     val compilationUnit: Property<SqlDelightCompilationUnit>
     val verifyMigrations: Property<Boolean>
     val verifyDefinitions: Property<Boolean>
+    val driverInitializer: Property<DriverInitializer>
   }
 
   abstract class VerifyMigrationAction : WorkAction<VerifyMigrationWorkParameters> {
@@ -116,6 +122,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
       if (!environment.dialect.isSqlite) return
       parameters.workingDirectory.get().asFile.deleteRecursively()
 
+      parameters.driverInitializer.orNull?.execute()
       val catalog = createCurrentDb()
 
       val databaseFiles = sourceFolders.asSequence()
@@ -199,3 +206,17 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
     }
   }
 }
+
+/**
+ * Allows consumers to configure and register (with [DriverManager]) their custom drivers prior to
+ * running migration verification task.
+ */
+interface DriverInitializer : Serializable {
+  fun execute()
+}
+
+class DriverInitializerImpl: DriverInitializer {
+  override fun execute() {
+    println("DriverInitializerImpl executed!")
+  }
+}

sqldelight-gradle-plugin/src/test/kotlin/app/cash/sqldelight/tests/MigrationTest.kt

@@ -36,6 +36,16 @@ class MigrationTest {
     assertThat(output.output).contains("BUILD SUCCESSFUL")
   }
 
+  @Test fun `driver initializer is executed`() {
+    val output = GradleRunner.create()
+            .withCommonConfiguration(File("src/test/migration-driver-initializer"))
+            .withArguments("clean", "verifyMainDatabaseMigration", "--stacktrace")
+            .build()
+
+    assertThat(output.output).contains("DriverInitializerImpl executed!")
+    assertThat(output.output).contains("BUILD SUCCESSFUL")
+  }
+
   @Test fun `failing migration errors properly`() {
     val output = GradleRunner.create()
       .withCommonConfiguration(File("src/test/migration-failure"))

sqldelight-gradle-plugin/src/test/migration-driver-initializer/build.gradle

@@ -0,0 +1,31 @@
+import app.cash.sqldelight.gradle.VerifyMigrationTask
+import app.cash.sqldelight.gradle.DriverInitializerImpl
+
+buildscript {
+  apply from: "${projectDir.absolutePath}/../buildscript.gradle"
+}
+
+apply plugin: 'org.jetbrains.kotlin.jvm'
+apply plugin: 'app.cash.sqldelight'
+
+sqldelight {
+  databases {
+    Database {
+      packageName = "com.example"
+      verifyMigrations = true
+    }
+  }
+}
+
+dependencies {
+  implementation libs.junit
+}
+
+tasks.whenTaskAdded {
+  if (it.name == "verifyMainDatabaseMigration") {
+    def task = it as VerifyMigrationTask
+    task.configure {
+      it.driverInitializer = new DriverInitializerImpl()
+    }
+  }
+}

sqldelight-gradle-plugin/src/test/migration-driver-initializer/settings.gradle

@@ -0,0 +1,3 @@
+apply from: "../settings.gradle"
+
+rootProject.name = 'sqldelight-migrations'

sqldelight-gradle-plugin/src/test/migration-driver-initializer/src/main/sqldelight/com/sample/Test.sq

@@ -0,0 +1,21 @@
+CREATE TABLE test(
+  value TEXT NOT NULL,
+  value2 TEXT
+);
+
+CREATE INDEX testIndex
+  ON test(value);
+
+CREATE TRIGGER testTrigger
+  AFTER DELETE ON test
+  BEGIN
+    INSERT INTO test VALUES ("1", "2");
+  END;
+
+CREATE VIEW testView AS
+  SELECT *
+  FROM test;
+
+select:
+SELECT *
+FROM test;

sqldelight-gradle-plugin/src/test/migration-driver-initializer/src/main/sqldelight/migrations/1.sqm

@@ -0,0 +1,13 @@
+ALTER TABLE test ADD COLUMN value2 TEXT;
+
+CREATE INDEX testIndex ON test(value);
+
+CREATE TRIGGER testTrigger
+AFTER DELETE ON test
+BEGIN
+INSERT INTO test VALUES ("1", "2");
+END;
+
+CREATE VIEW testView AS
+SELECT *
+FROM test;