Take Java migration classes into account in Flyway extension

extensions/flyway/deployment/src/main/java/io/quarkus/flyway/FlywayProcessor.java

@@ -3,6 +3,7 @@
 import static io.quarkus.deployment.annotations.ExecutionTime.STATIC_INIT;
 
 import java.io.IOException;
+import java.lang.reflect.Modifier;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
@@ -27,6 +28,9 @@
 import javax.enterprise.inject.Default;
 
 import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.migration.JavaMigration;
+import org.jboss.jandex.ClassInfo;
+import org.jboss.jandex.DotName;
 import org.jboss.logging.Logger;
 
 import io.quarkus.agroal.deployment.JdbcDataSourceBuildItem;
@@ -42,9 +46,13 @@
 import io.quarkus.deployment.annotations.Record;
 import io.quarkus.deployment.builditem.BytecodeTransformerBuildItem;
 import io.quarkus.deployment.builditem.CapabilityBuildItem;
+import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
+import io.quarkus.deployment.builditem.IndexDependencyBuildItem;
 import io.quarkus.deployment.builditem.ServiceStartBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.NativeImageResourceBuildItem;
+import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
+import io.quarkus.deployment.recording.RecorderContext;
 import io.quarkus.flyway.runtime.FlywayBuildTimeConfig;
 import io.quarkus.flyway.runtime.FlywayContainerProducer;
 import io.quarkus.flyway.runtime.FlywayRecorder;
@@ -57,6 +65,8 @@ class FlywayProcessor {
 
     private static final String FLYWAY_BEAN_NAME_PREFIX = "flyway_";
 
+    private static final DotName JAVA_MIGRATION = DotName.createSimple(JavaMigration.class.getName());
+
     private static final Logger LOGGER = Logger.getLogger(FlywayProcessor.class);
 
     FlywayBuildTimeConfig flywayBuildConfig;
@@ -73,11 +83,19 @@ void scannerTransformer(BuildProducer<BytecodeTransformerBuildItem> transformers
                         new ScannerTransformer()));
     }
 
+    @BuildStep
+    IndexDependencyBuildItem indexFlyway() {
+        return new IndexDependencyBuildItem("org.flywaydb", "flyway-core");
+    }
+
     @Record(STATIC_INIT)
     @BuildStep
     void build(BuildProducer<FeatureBuildItem> featureProducer,
             BuildProducer<NativeImageResourceBuildItem> resourceProducer,
+            BuildProducer<ReflectiveClassBuildItem> reflectiveClassProducer,
             FlywayRecorder recorder,
+            RecorderContext context,
+            CombinedIndexBuildItem combinedIndexBuildItem,
             List<JdbcDataSourceBuildItem> jdbcDataSourceBuildItems) throws IOException, URISyntaxException {
 
         featureProducer.produce(new FeatureBuildItem(FeatureBuildItem.FLYWAY));
@@ -87,9 +105,25 @@ void build(BuildProducer<FeatureBuildItem> featureProducer,
         List<String> applicationMigrations = discoverApplicationMigrations(getMigrationLocations(dataSourceNames));
         recorder.setApplicationMigrationFiles(applicationMigrations);
 
+        Set<Class<?>> javaMigrationClasses = new HashSet<>();
+        addJavaMigrations(combinedIndexBuildItem.getIndex().getAllKnownImplementors(JAVA_MIGRATION), context,
+                reflectiveClassProducer, javaMigrationClasses);
+        recorder.setApplicationMigrationClasses(javaMigrationClasses);
+
         resourceProducer.produce(new NativeImageResourceBuildItem(applicationMigrations.toArray(new String[0])));
     }
 
+    private void addJavaMigrations(Collection<ClassInfo> candidates, RecorderContext context,
+            BuildProducer<ReflectiveClassBuildItem> reflectiveClassProducer, Set<Class<?>> javaMigrationClasses) {
+        for (ClassInfo javaMigration : candidates) {
+            if (Modifier.isAbstract(javaMigration.flags())) {
+                continue;
+            }
+            javaMigrationClasses.add(context.classProxy(javaMigration.name().toString()));
+            reflectiveClassProducer.produce(new ReflectiveClassBuildItem(false, false, javaMigration.name().toString()));
+        }
+    }
+
     @BuildStep
     @Record(ExecutionTime.RUNTIME_INIT)
     ServiceStartBuildItem createBeansAndStartActions(FlywayRecorder recorder,

extensions/flyway/deployment/src/test/java/io/quarkus/flyway/test/FlywayExtensionCleanAndMigrateAtStartWithJavaMigrationTest.java

@@ -0,0 +1,101 @@
+package io.quarkus.flyway.test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import javax.inject.Inject;
+
+import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.MigrationVersion;
+import org.flywaydb.core.api.migration.BaseJavaMigration;
+import org.flywaydb.core.api.migration.Context;
+import org.flywaydb.core.api.migration.JavaMigration;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.agroal.api.AgroalDataSource;
+import io.quarkus.test.QuarkusUnitTest;
+
+public class FlywayExtensionCleanAndMigrateAtStartWithJavaMigrationTest {
+
+    @Inject
+    Flyway flyway;
+
+    @Inject
+    AgroalDataSource defaultDataSource;
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest()
+            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class)
+                    .addClasses(V1_0_1__Update.class, V1_0_2__Update.class)
+                    .addAsResource("db/migration/V1.0.0__Quarkus.sql")
+                    .addAsResource("clean-and-migrate-at-start-config.properties", "application.properties"));
+
+    @Test
+    @DisplayName("Clean and migrate at start correctly")
+    public void testFlywayConfigInjection() throws SQLException {
+
+        try (Connection connection = defaultDataSource.getConnection(); Statement stat = connection.createStatement()) {
+            try (ResultSet countQuery = stat.executeQuery("select count(1) from quarked_flyway")) {
+                assertTrue(countQuery.first());
+                assertEquals(2,
+                        countQuery.getInt(1),
+                        "Table 'quarked_flyway' does not contain the expected number of rows");
+            }
+        }
+        String currentVersion = flyway.info().current().getVersion().toString();
+        assertEquals("1.0.2", currentVersion, "Expected to be 1.0.2 as there is a SQL and two Java migration scripts");
+    }
+
+    public static class V1_0_1__Update extends BaseJavaMigration {
+        @Override
+        public void migrate(Context context) throws Exception {
+            try (Statement statement = context.getConnection().createStatement()) {
+                statement.executeUpdate("INSERT INTO quarked_flyway VALUES (1001, 'test')");
+            }
+        }
+    }
+
+    public static class V1_0_2__Update implements JavaMigration {
+        @Override
+        public MigrationVersion getVersion() {
+            return MigrationVersion.fromVersion("1.0.2");
+        }
+
+        @Override
+        public String getDescription() {
+            return getClass().getSimpleName();
+        }
+
+        @Override
+        public Integer getChecksum() {
+            return null;
+        }
+
+        @Override
+        public boolean isUndo() {
+            return false;
+        }
+
+        @Override
+        public boolean canExecuteInTransaction() {
+            return true;
+        }
+
+        @Override
+        public void migrate(Context context) throws Exception {
+            try (Statement statement = context.getConnection().createStatement()) {
+                statement.executeUpdate("INSERT INTO quarked_flyway VALUES (1002, 'test')");
+            }
+        }
+    }
+
+}

extensions/flyway/runtime/src/main/java/io/quarkus/flyway/runtime/FlywayRecorder.java

@@ -1,6 +1,7 @@
 package io.quarkus.flyway.runtime;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.function.Supplier;
 
@@ -20,11 +21,16 @@ public class FlywayRecorder {
 
     private final List<FlywayContainer> flywayContainers = new ArrayList<>(2);
 
-    public void setApplicationMigrationFiles(List<String> migrationFiles) {
+    public void setApplicationMigrationFiles(Collection<String> migrationFiles) {
         log.debugv("Setting the following application migration files: {0}", migrationFiles);
         QuarkusPathLocationScanner.setApplicationMigrationFiles(migrationFiles);
     }
 
+    public void setApplicationMigrationClasses(Collection<Class<?>> migrationClasses) {
+        log.debugv("Setting the following application migration classes: {0}", migrationClasses);
+        QuarkusPathLocationScanner.setApplicationMigrationClasses(migrationClasses);
+    }
+
     public Supplier<Flyway> flywaySupplier(String dataSourceName) {
         DataSource dataSource = DataSources.fromName(dataSourceName);
         FlywayContainerProducer flywayProducer = Arc.container().instance(FlywayContainerProducer.class).get();

extensions/flyway/runtime/src/main/java/io/quarkus/flyway/runtime/QuarkusPathLocationScanner.java

@@ -3,8 +3,6 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
 
 import org.flywaydb.core.api.Location;
 import org.flywaydb.core.internal.resource.LoadableResource;
@@ -20,7 +18,8 @@
 public final class QuarkusPathLocationScanner implements ResourceAndClassScanner {
     private static final Logger LOGGER = Logger.getLogger(QuarkusPathLocationScanner.class);
     private static final String LOCATION_SEPARATOR = "/";
-    private static List<String> applicationMigrationFiles;
+    private static Collection<String> applicationMigrationFiles;
+    private static Collection<Class<?>> applicationMigrationClasses;
 
     private final Collection<LoadableResource> scannedResources;
 
@@ -74,11 +73,14 @@ private boolean canHandleMigrationFile(Collection<Location> locations, String mi
      */
     @Override
     public Collection<Class<?>> scanForClasses() {
-        // Classes are not supported in native mode
-        return Collections.emptyList();
+        return applicationMigrationClasses;
     }
 
-    public static void setApplicationMigrationFiles(List<String> applicationMigrationFiles) {
+    public static void setApplicationMigrationFiles(Collection<String> applicationMigrationFiles) {
         QuarkusPathLocationScanner.applicationMigrationFiles = applicationMigrationFiles;
     }
+
+    public static void setApplicationMigrationClasses(Collection<Class<?>> applicationMigrationClasses) {
+        QuarkusPathLocationScanner.applicationMigrationClasses = applicationMigrationClasses;
+    }
 }