Allow registering DriverInitializer for VerifyMigrationTask with ServiceLoader mechanism

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

@@ -13,6 +13,7 @@ import org.gradle.api.Project
 import org.gradle.api.file.ConfigurableFileCollection
 import org.gradle.api.file.DirectoryProperty
 import org.gradle.api.internal.catalog.DelegatingProjectDependency
+import org.gradle.api.provider.MapProperty
 import org.gradle.api.provider.Property
 import org.gradle.api.provider.Provider
 import java.io.File
@@ -32,6 +33,7 @@ abstract class SqlDelightDatabase @Inject constructor(
   abstract val migrationOutputDirectory: DirectoryProperty
   val migrationOutputFileFormat: Property<String> = project.objects.property(String::class.java).convention(".sql")
   val generateAsync: Property<Boolean> = project.objects.property(Boolean::class.java).convention(false)
+  val connectionProperties: MapProperty<String, String> = project.objects.mapProperty(String::class.java, String::class.java)
 
   internal val configuration = project.configurations.create("${name}DialectClasspath").apply {
     isCanBeConsumed = false
@@ -263,6 +265,7 @@ abstract class SqlDelightDatabase @Inject constructor(
         it.verifyMigrations = verifyMigrations.get()
         it.classpath.setFrom(configuration.fileCollection { true })
         it.classpath.from(moduleConfiguration.fileCollection { true })
+        it.connectionProperties.set(connectionProperties)
       }
 
     if (schemaOutputDirectory.getOrNull() != null) {

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

@@ -14,6 +14,7 @@ import app.cash.sqlite.migrations.findDatabaseFiles
 import org.gradle.api.file.DirectoryProperty
 import org.gradle.api.file.FileTree
 import org.gradle.api.logging.Logging
+import org.gradle.api.provider.MapProperty
 import org.gradle.api.provider.Property
 import org.gradle.api.tasks.CacheableTask
 import org.gradle.api.tasks.IgnoreEmptyDirectories
@@ -29,8 +30,11 @@ import org.gradle.api.tasks.TaskAction
 import org.gradle.workers.WorkAction
 import org.gradle.workers.WorkParameters
 import java.io.File
+import java.util.Properties
 import java.util.ServiceLoader
 
+interface MigrationVerificationDriver
+
 @CacheableTask
 abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
   @Suppress("unused")
@@ -51,6 +55,8 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
 
   @Input var verifyDefinitions: Boolean = true
 
+  @get:Input abstract val connectionProperties: MapProperty<String, String>
+
   /* 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 +74,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
         it.verifyMigrations.set(verifyMigrations)
         it.compilationUnit.set(compilationUnit)
         it.verifyDefinitions.set(verifyDefinitions)
+        it.connectionProperties.set(connectionProperties.get())
       }
       workQueue.await()
     }.onSuccess {
@@ -92,6 +99,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
     val compilationUnit: Property<SqlDelightCompilationUnit>
     val verifyMigrations: Property<Boolean>
     val verifyDefinitions: Property<Boolean>
+    val connectionProperties: MapProperty<String, String>
   }
 
   abstract class VerifyMigrationAction : WorkAction<VerifyMigrationWorkParameters> {
@@ -116,7 +124,13 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
       if (!environment.dialect.isSqlite) return
       parameters.workingDirectory.get().asFile.deleteRecursively()
 
-      val catalog = createCurrentDb()
+      val connectionPropertiesMap = parameters.connectionProperties.get()
+      val connectionProperties = Properties()
+      connectionPropertiesMap.forEach { (key, value) ->
+        connectionProperties[key] = value
+      }
+      ServiceLoader.load(MigrationVerificationDriver::class.java).findFirst()
+      val catalog = createCurrentDb(connectionProperties)
 
       val databaseFiles = sourceFolders.asSequence()
         .findDatabaseFiles()
@@ -126,13 +140,13 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
       }
 
       databaseFiles.forEach { dbFile ->
-        checkMigration(dbFile, catalog)
+        checkMigration(dbFile, catalog, connectionProperties)
       }
 
       checkForGaps()
     }
 
-    private fun createCurrentDb(): CatalogDatabase {
+    private fun createCurrentDb(connectionProperties: Properties): CatalogDatabase {
       val sourceFiles = ArrayList<SqlDelightQueriesFile>()
       environment.forSourceFiles { file ->
         if (file is SqlDelightQueriesFile) sourceFiles.add(file)
@@ -143,11 +157,11 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
       ) { sqlText ->
         initStatements.add(CatalogDatabase.InitStatement(sqlText, "Error compiling $sqlText"))
       }
-      return CatalogDatabase.withInitStatements(initStatements)
+      return CatalogDatabase.withInitStatements(initStatements, connectionProperties)
     }
 
-    private fun checkMigration(dbFile: File, currentDb: CatalogDatabase) {
-      val actualCatalog = createActualDb(dbFile)
+    private fun checkMigration(dbFile: File, currentDb: CatalogDatabase, connectionProperties: Properties) {
+      val actualCatalog = createActualDb(dbFile, connectionProperties)
       val databaseComparator = ObjectDifferDatabaseComparator(
         ignoreDefinitions = !parameters.verifyDefinitions.get(),
         circularReferenceExceptionLogger = {
@@ -164,7 +178,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
       }
     }
 
-    private fun createActualDb(dbFile: File): CatalogDatabase {
+    private fun createActualDb(dbFile: File, connectionProperties: Properties): CatalogDatabase {
       val version = dbFile.nameWithoutExtension.toInt()
       val copy = dbFile.copyTo(File(parameters.workingDirectory.get().asFile, dbFile.name))
       val initStatements = ArrayList<CatalogDatabase.InitStatement>()
@@ -182,7 +196,7 @@ abstract class VerifyMigrationTask : SqlDelightWorkerTask() {
           )
         }
       }
-      return CatalogDatabase.fromFile(copy.absolutePath, initStatements).also { copy.delete() }
+      return CatalogDatabase.fromFile(copy.absolutePath, initStatements, connectionProperties).also { copy.delete() }
     }
 
     private fun checkForGaps() {

sqlite-migrations/src/main/kotlin/app/cash/sqlite/migrations/CatalogDatabase.kt

@@ -9,6 +9,7 @@ import schemacrawler.tools.utility.SchemaCrawlerUtility
 import java.sql.Connection
 import java.sql.DriverManager
 import java.sql.SQLException
+import java.util.Properties
 
 class CatalogDatabase private constructor(
   internal val catalog: Catalog,
@@ -28,21 +29,28 @@ class CatalogDatabase private constructor(
           .toOptions(),
       )
 
-    fun withInitStatements(initStatements: List<InitStatement>): CatalogDatabase {
-      return fromFile("", initStatements)
+    fun withInitStatements(
+      initStatements: List<InitStatement>,
+      connectionProperties: Properties,
+    ): CatalogDatabase {
+      return fromFile("", initStatements, connectionProperties)
     }
 
-    fun fromFile(path: String, initStatements: List<InitStatement>): CatalogDatabase {
-      return createConnection(path).init(initStatements).use {
+    fun fromFile(
+      path: String,
+      initStatements: List<InitStatement>,
+      connectionProperties: Properties,
+    ): CatalogDatabase {
+      return createConnection(path, connectionProperties).init(initStatements).use {
         CatalogDatabase(SchemaCrawlerUtility.getCatalog(it, schemaCrawlerOptions))
       }
     }
 
-    private fun createConnection(path: String): Connection {
+    private fun createConnection(path: String, connectionProperties: Properties): Connection {
       return try {
-        DriverManager.getConnection("jdbc:sqlite:$path")
+        DriverManager.getConnection("jdbc:sqlite:$path", connectionProperties)
       } catch (e: SQLException) {
-        DriverManager.getConnection("jdbc:sqlite:$path")
+        DriverManager.getConnection("jdbc:sqlite:$path", connectionProperties)
       }
     }