Merge pull request #35315 from geoand/#35314

Fix Datasource timing issues with Liquibase / Flyway and OpenTelemetry

extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSources.java

@@ -113,7 +113,10 @@ public DataSources(DataSourcesBuildTimeConfig dataSourcesBuildTimeConfig,
      * (which makes sense because {@code DataSource} is a {@code Singleton} bean).
      * <p>
      * This method is thread-safe
+     *
+     * @deprecated This method should not be used as it can very easily lead to timing issues during bean creation
      */
+    @Deprecated
     public static AgroalDataSource fromName(String dataSourceName) {
         return Arc.container().instance(DataSources.class).get()
                 .getDataSource(dataSourceName);

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

@@ -18,18 +18,21 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import jakarta.enterprise.context.Dependent;
 import jakarta.enterprise.inject.Default;
+import jakarta.inject.Singleton;
 
 import org.flywaydb.core.Flyway;
 import org.flywaydb.core.api.Location;
 import org.flywaydb.core.api.callback.Callback;
 import org.flywaydb.core.api.migration.JavaMigration;
 import org.flywaydb.core.extensibility.Plugin;
+import org.jboss.jandex.AnnotationInstance;
 import org.jboss.jandex.ClassInfo;
+import org.jboss.jandex.ClassType;
 import org.jboss.jandex.DotName;
 import org.jboss.logging.Logger;
 
+import io.quarkus.agroal.runtime.DataSources;
 import io.quarkus.agroal.spi.JdbcDataSourceBuildItem;
 import io.quarkus.agroal.spi.JdbcDataSourceSchemaReadyBuildItem;
 import io.quarkus.agroal.spi.JdbcInitialSQLGeneratorBuildItem;
@@ -62,6 +65,7 @@
 import io.quarkus.deployment.logging.LoggingSetupBuildItem;
 import io.quarkus.deployment.recording.RecorderContext;
 import io.quarkus.flyway.runtime.FlywayBuildTimeConfig;
+import io.quarkus.flyway.runtime.FlywayContainer;
 import io.quarkus.flyway.runtime.FlywayContainerProducer;
 import io.quarkus.flyway.runtime.FlywayRecorder;
 import io.quarkus.flyway.runtime.FlywayRuntimeConfig;
@@ -71,6 +75,7 @@ class FlywayProcessor {
 
     private static final String CLASSPATH_APPLICATION_MIGRATIONS_PROTOCOL = "classpath";
 
+    private static final String FLYWAY_CONTAINER_BEAN_NAME_PREFIX = "flyway_container_";
     private static final String FLYWAY_BEAN_NAME_PREFIX = "flyway_";
 
     private static final DotName JAVA_MIGRATION = DotName.createSimple(JavaMigration.class.getName());
@@ -172,8 +177,6 @@ void createBeans(FlywayRecorder recorder,
         // add the @FlywayDataSource class otherwise it won't be registered as a qualifier
         additionalBeans.produce(AdditionalBeanBuildItem.builder().addBeanClass(FlywayDataSource.class).build());
 
-        recorder.resetFlywayContainers();
-
         Collection<String> dataSourceNames = getDataSourceNames(jdbcDataSourceBuildItems);
 
         for (String dataSourceName : dataSourceNames) {
@@ -182,25 +185,62 @@ void createBeans(FlywayRecorder recorder,
             if (!hasMigrations) {
                 createPossible = sqlGeneratorBuildItems.stream().anyMatch(s -> s.getDatabaseName().equals(dataSourceName));
             }
-            SyntheticBeanBuildItem.ExtendedBeanConfigurator configurator = SyntheticBeanBuildItem
+
+            SyntheticBeanBuildItem.ExtendedBeanConfigurator flywayContainerConfigurator = SyntheticBeanBuildItem
+                    .configure(FlywayContainer.class)
+                    .scope(Singleton.class)
+                    .setRuntimeInit()
+                    .unremovable()
+                    .addInjectionPoint(ClassType.create(DotName.createSimple(FlywayContainerProducer.class)))
+                    .addInjectionPoint(ClassType.create(DotName.createSimple(DataSources.class)))
+                    .createWith(recorder.flywayContainerFunction(dataSourceName, hasMigrations, createPossible));
+
+            AnnotationInstance flywayContainerQualifier;
+
+            if (DataSourceUtil.isDefault(dataSourceName)) {
+                flywayContainerConfigurator.addQualifier(Default.class);
+
+                // Flyway containers used to be ordered with the default database coming first.
+                // Some multitenant tests are relying on this order.
+                flywayContainerConfigurator.priority(10);
+
+                flywayContainerQualifier = AnnotationInstance.builder(Default.class).build();
+            } else {
+                String beanName = FLYWAY_CONTAINER_BEAN_NAME_PREFIX + dataSourceName;
+                flywayContainerConfigurator.name(beanName);
+
+                flywayContainerConfigurator.addQualifier().annotation(DotNames.NAMED).addValue("value", beanName).done();
+                flywayContainerConfigurator.addQualifier().annotation(FlywayDataSource.class).addValue("value", dataSourceName)
+                        .done();
+                flywayContainerConfigurator.priority(5);
+
+                flywayContainerQualifier = AnnotationInstance.builder(FlywayDataSource.class).add("value", dataSourceName)
+                        .build();
+            }
+
+            syntheticBeanBuildItemBuildProducer.produce(flywayContainerConfigurator.done());
+
+            SyntheticBeanBuildItem.ExtendedBeanConfigurator flywayConfigurator = SyntheticBeanBuildItem
                     .configure(Flyway.class)
-                    .scope(Dependent.class) // this is what the existing code does, but it doesn't seem reasonable
+                    .scope(Singleton.class)
                     .setRuntimeInit()
                     .unremovable()
-                    .supplier(recorder.flywaySupplier(dataSourceName,
-                            hasMigrations, createPossible));
+                    .addInjectionPoint(ClassType.create(DotName.createSimple(FlywayContainer.class)), flywayContainerQualifier)
+                    .createWith(recorder.flywayFunction(dataSourceName));
 
             if (DataSourceUtil.isDefault(dataSourceName)) {
-                configurator.addQualifier(Default.class);
+                flywayConfigurator.addQualifier(Default.class);
+                flywayConfigurator.priority(10);
             } else {
                 String beanName = FLYWAY_BEAN_NAME_PREFIX + dataSourceName;
-                configurator.name(beanName);
+                flywayConfigurator.name(beanName);
+                flywayConfigurator.priority(5);
 
-                configurator.addQualifier().annotation(DotNames.NAMED).addValue("value", beanName).done();
-                configurator.addQualifier().annotation(FlywayDataSource.class).addValue("value", dataSourceName).done();
+                flywayConfigurator.addQualifier().annotation(DotNames.NAMED).addValue("value", beanName).done();
+                flywayConfigurator.addQualifier().annotation(FlywayDataSource.class).addValue("value", dataSourceName).done();
             }
 
-            syntheticBeanBuildItemBuildProducer.produce(configurator.done());
+            syntheticBeanBuildItemBuildProducer.produce(flywayConfigurator.done());
         }
     }
 

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

@@ -1,24 +1,30 @@
 package io.quarkus.flyway.runtime;
 
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Comparator;
+import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.function.Supplier;
 
+import io.quarkus.arc.Arc;
+import io.quarkus.arc.InstanceHandle;
 import io.quarkus.datasource.common.runtime.DataSourceUtil;
 
 public class FlywayContainersSupplier implements Supplier<Collection<FlywayContainer>> {
 
     @Override
     public Collection<FlywayContainer> get() {
-        if (FlywayRecorder.FLYWAY_CONTAINERS.isEmpty()) {
-            return Collections.emptySet();
+        List<InstanceHandle<FlywayContainer>> flywayContainerHandles = Arc.container().listAll(FlywayContainer.class);
+
+        if (flywayContainerHandles.isEmpty()) {
+            return Set.of();
         }
 
         Set<FlywayContainer> containers = new TreeSet<>(FlywayContainerComparator.INSTANCE);
-        containers.addAll(FlywayRecorder.FLYWAY_CONTAINERS);
+        for (InstanceHandle<FlywayContainer> flywayContainerHandle : flywayContainerHandles) {
+            containers.add(flywayContainerHandle.get());
+        }
         return containers;
     }
 

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

@@ -1,13 +1,13 @@
 package io.quarkus.flyway.runtime;
 
-import java.util.ArrayList;
+import java.lang.annotation.Annotation;
 import java.util.Collection;
-import java.util.List;
 import java.util.Map;
-import java.util.function.Supplier;
+import java.util.function.Function;
 
 import javax.sql.DataSource;
 
+import jakarta.enterprise.inject.Default;
 import jakarta.enterprise.inject.UnsatisfiedResolutionException;
 
 import org.flywaydb.core.Flyway;
@@ -18,6 +18,10 @@
 import io.quarkus.agroal.runtime.DataSources;
 import io.quarkus.agroal.runtime.UnconfiguredDataSource;
 import io.quarkus.arc.Arc;
+import io.quarkus.arc.InstanceHandle;
+import io.quarkus.arc.SyntheticCreationalContext;
+import io.quarkus.datasource.common.runtime.DataSourceUtil;
+import io.quarkus.flyway.FlywayDataSource.FlywayDataSourceLiteral;
 import io.quarkus.runtime.RuntimeValue;
 import io.quarkus.runtime.annotations.Recorder;
 
@@ -26,8 +30,6 @@ public class FlywayRecorder {
 
     private static final Logger log = Logger.getLogger(FlywayRecorder.class);
 
-    static final List<FlywayContainer> FLYWAY_CONTAINERS = new ArrayList<>(2);
-
     private final RuntimeValue<FlywayRuntimeConfig> config;
 
     public FlywayRecorder(RuntimeValue<FlywayRuntimeConfig> config) {
@@ -49,27 +51,37 @@ public void setApplicationCallbackClasses(Map<String, Collection<Callback>> call
         QuarkusPathLocationScanner.setApplicationCallbackClasses(callbackClasses);
     }
 
-    public void resetFlywayContainers() {
-        FLYWAY_CONTAINERS.clear();
-    }
-
-    public Supplier<Flyway> flywaySupplier(String dataSourceName, boolean hasMigrations, boolean createPossible) {
-        DataSource dataSource = DataSources.fromName(dataSourceName);
-        if (dataSource instanceof UnconfiguredDataSource) {
-            return new Supplier<Flyway>() {
-                @Override
-                public Flyway get() {
+    public Function<SyntheticCreationalContext<FlywayContainer>, FlywayContainer> flywayContainerFunction(String dataSourceName,
+            boolean hasMigrations,
+            boolean createPossible) {
+        return new Function<>() {
+            @Override
+            public FlywayContainer apply(SyntheticCreationalContext<FlywayContainer> context) {
+                DataSource dataSource = context.getInjectedReference(DataSources.class).getDataSource(dataSourceName);
+                if (dataSource instanceof UnconfiguredDataSource) {
                     throw new UnsatisfiedResolutionException("No datasource present");
                 }
-            };
-        }
-        FlywayContainerProducer flywayProducer = Arc.container().instance(FlywayContainerProducer.class).get();
-        FlywayContainer flywayContainer = flywayProducer.createFlyway(dataSource, dataSourceName, hasMigrations,
-                createPossible);
-        FLYWAY_CONTAINERS.add(flywayContainer);
-        return new Supplier<Flyway>() {
+
+                FlywayContainerProducer flywayProducer = context.getInjectedReference(FlywayContainerProducer.class);
+                FlywayContainer flywayContainer = flywayProducer.createFlyway(dataSource, dataSourceName, hasMigrations,
+                        createPossible);
+                return flywayContainer;
+            }
+        };
+    }
+
+    public Function<SyntheticCreationalContext<Flyway>, Flyway> flywayFunction(String dataSourceName) {
+        return new Function<>() {
             @Override
-            public Flyway get() {
+            public Flyway apply(SyntheticCreationalContext<Flyway> context) {
+                Annotation flywayContainerQualifier;
+                if (DataSourceUtil.isDefault(dataSourceName)) {
+                    flywayContainerQualifier = Default.Literal.INSTANCE;
+                } else {
+                    flywayContainerQualifier = FlywayDataSourceLiteral.of(dataSourceName);
+                }
+
+                FlywayContainer flywayContainer = context.getInjectedReference(FlywayContainer.class, flywayContainerQualifier);
                 return flywayContainer.getFlyway();
             }
         };
@@ -79,7 +91,10 @@ public void doStartActions() {
         if (!config.getValue().enabled) {
             return;
         }
-        for (FlywayContainer flywayContainer : FLYWAY_CONTAINERS) {
+
+        for (InstanceHandle<FlywayContainer> flywayContainerHandle : Arc.container().listAll(FlywayContainer.class)) {
+            FlywayContainer flywayContainer = flywayContainerHandle.get();
+
             if (flywayContainer.isCleanAtStart()) {
                 flywayContainer.getFlyway().clean();
             }

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

@@ -1,23 +1,25 @@
 package io.quarkus.flyway.runtime;
 
+import io.quarkus.arc.Arc;
 import io.quarkus.datasource.runtime.DatabaseSchemaProvider;
 
 public class FlywaySchemaProvider implements DatabaseSchemaProvider {
+
     @Override
     public void resetDatabase(String dbName) {
-        for (FlywayContainer i : FlywayRecorder.FLYWAY_CONTAINERS) {
-            if (i.getDataSourceName().equals(dbName)) {
-                i.getFlyway().clean();
-                i.getFlyway().migrate();
+        for (FlywayContainer flywayContainer : Arc.container().select(FlywayContainer.class)) {
+            if (flywayContainer.getDataSourceName().equals(dbName)) {
+                flywayContainer.getFlyway().clean();
+                flywayContainer.getFlyway().migrate();
             }
         }
     }
 
     @Override
     public void resetAllDatabases() {
-        for (FlywayContainer i : FlywayRecorder.FLYWAY_CONTAINERS) {
-            i.getFlyway().clean();
-            i.getFlyway().migrate();
+        for (FlywayContainer flywayContainer : Arc.container().select(FlywayContainer.class)) {
+            flywayContainer.getFlyway().clean();
+            flywayContainer.getFlyway().migrate();
         }
     }
 }

extensions/liquibase/deployment/src/main/java/io/quarkus/liquibase/deployment/LiquibaseProcessor.java

@@ -28,6 +28,7 @@
 import org.jboss.jandex.DotName;
 import org.jboss.logging.Logger;
 
+import io.quarkus.agroal.runtime.DataSources;
 import io.quarkus.agroal.spi.JdbcDataSourceBuildItem;
 import io.quarkus.agroal.spi.JdbcDataSourceSchemaReadyBuildItem;
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
@@ -285,6 +286,7 @@ void createBeans(LiquibaseRecorder recorder,
                     .setRuntimeInit()
                     .unremovable()
                     .addInjectionPoint(ClassType.create(DotName.createSimple(LiquibaseFactoryProducer.class)))
+                    .addInjectionPoint(ClassType.create(DotName.createSimple(DataSources.class)))
                     .createWith(recorder.liquibaseFunction(dataSourceName));
 
             if (DataSourceUtil.isDefault(dataSourceName)) {

extensions/liquibase/runtime/src/main/java/io/quarkus/liquibase/runtime/LiquibaseRecorder.java

@@ -29,21 +29,16 @@ public LiquibaseRecorder(RuntimeValue<LiquibaseRuntimeConfig> config) {
     }
 
     public Function<SyntheticCreationalContext<LiquibaseFactory>, LiquibaseFactory> liquibaseFunction(String dataSourceName) {
-        DataSource dataSource = DataSources.fromName(dataSourceName);
-        if (dataSource instanceof UnconfiguredDataSource) {
-            return new Function<SyntheticCreationalContext<LiquibaseFactory>, LiquibaseFactory>() {
-                @Override
-                public LiquibaseFactory apply(SyntheticCreationalContext<LiquibaseFactory> context) {
-                    throw new UnsatisfiedResolutionException("No datasource has been configured");
-                }
-            };
-        }
         return new Function<SyntheticCreationalContext<LiquibaseFactory>, LiquibaseFactory>() {
             @Override
             public LiquibaseFactory apply(SyntheticCreationalContext<LiquibaseFactory> context) {
+                DataSource dataSource = context.getInjectedReference(DataSources.class).getDataSource(dataSourceName);
+                if (dataSource instanceof UnconfiguredDataSource) {
+                    throw new UnsatisfiedResolutionException("No datasource has been configured");
+                }
+
                 LiquibaseFactoryProducer liquibaseProducer = context.getInjectedReference(LiquibaseFactoryProducer.class);
-                LiquibaseFactory liquibaseFactory = liquibaseProducer.createLiquibaseFactory(dataSource, dataSourceName);
-                return liquibaseFactory;
+                return liquibaseProducer.createLiquibaseFactory(dataSource, dataSourceName);
             }
         };
     }