Support multiple Hibernate ORM persistence units via Quarkus configuration

extensions/hibernate-orm/deployment/pom.xml

@@ -68,6 +68,11 @@
             <artifactId>rest-assured</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-smallrye-metrics-deployment</artifactId>

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateOrmConfig.java

@@ -1,7 +1,9 @@
 package io.quarkus.hibernate.orm.deployment;
 
+import java.util.Map;
 import java.util.Optional;
 
+import io.quarkus.runtime.annotations.ConfigDocMapKey;
 import io.quarkus.runtime.annotations.ConfigDocSection;
 import io.quarkus.runtime.annotations.ConfigGroup;
 import io.quarkus.runtime.annotations.ConfigItem;
@@ -16,6 +18,16 @@ public class HibernateOrmConfig {
     @ConfigItem(name = ConfigItem.PARENT)
     public HibernateOrmConfigPersistenceUnit defaultPersistenceUnit;
 
+    /**
+     * Additional named persistence units.
+     */
+    @ConfigDocSection
+    @ConfigDocMapKey("persistence-unit-name")
+    // @ConfigItem(name = ConfigItem.PARENT)
+    // TODO MULTI-PUS for now, we keep the `persistence-units` but we will need to discuss it
+    // in most extensions, we decided to drop that part
+    public Map<String, HibernateOrmConfigPersistenceUnit> persistenceUnits;
+
     /**
      * Logging configuration.
      */
@@ -58,6 +70,7 @@ public class HibernateOrmConfig {
 
     public boolean isAnyPropertySet() {
         return defaultPersistenceUnit.isAnyPropertySet() ||
+                !persistenceUnits.isEmpty() ||
                 log.isAnyPropertySet() ||
                 statistics.isPresent() ||
                 metricsEnabled ||

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateOrmConfigPersistenceUnit.java

@@ -6,6 +6,7 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
+import java.util.Set;
 
 import io.quarkus.runtime.annotations.ConfigDocSection;
 import io.quarkus.runtime.annotations.ConfigGroup;
@@ -21,6 +22,11 @@ public class HibernateOrmConfigPersistenceUnit {
      */
     public Optional<String> datasource;
 
+    /**
+     * The packages in which the entities affected to this persistence unit are located.
+     */
+    public Optional<Set<String>> packages;
+
     /**
      * Class name of the Hibernate ORM dialect. The complete list of bundled dialects is available in the
      * https://docs.jboss.org/hibernate/stable/orm/javadocs/org/hibernate/dialect/package-summary.html[Hibernate ORM JavaDoc].

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateOrmProcessor.java

@@ -55,6 +55,7 @@
 import org.jboss.jandex.DotName;
 import org.jboss.jandex.IndexView;
 import org.jboss.jandex.Indexer;
+import org.jboss.logging.Logger;
 import org.jboss.logmanager.Level;
 
 import io.quarkus.agroal.spi.JdbcDataSourceBuildItem;
@@ -125,6 +126,9 @@ public final class HibernateOrmProcessor {
 
     public static final String HIBERNATE_ORM_CONFIG_PREFIX = "quarkus.hibernate-orm.";
     public static final String NO_SQL_LOAD_SCRIPT_FILE = "no-file";
+    public static final String DEFAULT_PERSISTENCE_UNIT_NAME = "default";
+
+    private static final Logger LOG = Logger.getLogger(HibernateOrmProcessor.class);
 
     private static final DotName PERSISTENCE_CONTEXT = DotName.createSimple(PersistenceContext.class.getName());
     private static final DotName PERSISTENCE_UNIT = DotName.createSimple(PersistenceUnit.class.getName());
@@ -212,14 +216,14 @@ public void configurationDescriptorBuilding(
             List<JdbcDataSourceBuildItem> jdbcDataSources,
             ApplicationArchivesBuildItem applicationArchivesBuildItem,
             LaunchModeBuildItem launchMode,
-            JpaEntitiesBuildItem domainObjects,
+            JpaEntitiesBuildItem jpaEntities,
             List<NonJpaModelBuildItem> nonJpaModelBuildItems,
             BuildProducer<SystemPropertyBuildItem> systemProperties,
             BuildProducer<NativeImageResourceBuildItem> nativeImageResources,
             BuildProducer<HotDeploymentWatchedFileBuildItem> hotDeploymentWatchedFiles,
             BuildProducer<PersistenceUnitDescriptorBuildItem> persistenceUnitDescriptors) {
 
-        if (!hasEntities(domainObjects, nonJpaModelBuildItems)) {
+        if (!hasEntities(jpaEntities, nonJpaModelBuildItems)) {
             // we can bail out early as there are no entities
             return;
         }
@@ -233,7 +237,7 @@ public void configurationDescriptorBuilding(
 
         if (impliedPU.shouldGenerateImpliedBlockingPersistenceUnit()) {
             handleHibernateORMWithNoPersistenceXml(hibernateOrmConfig, persistenceXmlDescriptors,
-                    jdbcDataSources, applicationArchivesBuildItem, launchMode.getLaunchMode(),
+                    jdbcDataSources, applicationArchivesBuildItem, launchMode.getLaunchMode(), jpaEntities,
                     systemProperties, nativeImageResources, hotDeploymentWatchedFiles, persistenceUnitDescriptors);
         }
     }
@@ -497,9 +501,9 @@ void registerBeans(HibernateOrmConfig hibernateOrmConfig,
                 .addBeanClasses(unremovableClasses.toArray(new Class<?>[unremovableClasses.size()]))
                 .build());
 
-        if (descriptors.size() == 1) {
-            // There is only one persistence unit - register CDI beans for EM and EMF if no
-            // producers are defined
+        if (descriptors.size() == 1 && hibernateOrmConfig.persistenceUnits.isEmpty()) {
+            // There is only one persistence unit, either via persistence.xml or it is the default one
+            // Register CDI beans for EM and EMF if no producers are defined
             if (isUserDefinedProducerMissing(combinedIndex.getIndex(), PERSISTENCE_UNIT)) {
                 additionalBeans.produce(new AdditionalBeanBuildItem(DefaultEntityManagerFactoryProducer.class));
             }
@@ -600,12 +604,13 @@ private void handleHibernateORMWithNoPersistenceXml(
             List<JdbcDataSourceBuildItem> jdbcDataSources,
             ApplicationArchivesBuildItem applicationArchivesBuildItem,
             LaunchMode launchMode,
+            JpaEntitiesBuildItem jpaEntities,
             BuildProducer<SystemPropertyBuildItem> systemProperties,
             BuildProducer<NativeImageResourceBuildItem> nativeImageResources,
             BuildProducer<HotDeploymentWatchedFileBuildItem> hotDeploymentWatchedFiles,
             BuildProducer<PersistenceUnitDescriptorBuildItem> persistenceUnitDescriptors) {
         if (!descriptors.isEmpty()) {
-            if (hibernateOrmConfig.isAnyPropertySet()) {
+            if (hibernateOrmConfig.isAnyPropertySet() || !hibernateOrmConfig.persistenceUnits.isEmpty()) {
                 throw new ConfigurationError(
                         "Hibernate ORM configuration present in persistence.xml and Quarkus config file at the same time\n"
                                 + "If you use persistence.xml remove all " + HIBERNATE_ORM_CONFIG_PREFIX
@@ -615,24 +620,44 @@ private void handleHibernateORMWithNoPersistenceXml(
             }
         }
 
-        producePersistenceUnitDescriptorFromConfig(
-                hibernateOrmConfig, "default", hibernateOrmConfig.defaultPersistenceUnit, jdbcDataSources,
-                applicationArchivesBuildItem, launchMode,
-                systemProperties, nativeImageResources, hotDeploymentWatchedFiles, persistenceUnitDescriptors);
+        Map<String, Set<String>> modelClassesPerPersistencesUnits = getModelClassesPerPersistenceUnits(hibernateOrmConfig,
+                jpaEntities);
+
+        Optional<JdbcDataSourceBuildItem> defaultJdbcDataSource = jdbcDataSources.stream()
+                .filter(i -> i.isDefault())
+                .findFirst();
+
+        if ((defaultJdbcDataSource.isPresent() && hibernateOrmConfig.persistenceUnits.isEmpty()) ||
+                hibernateOrmConfig.defaultPersistenceUnit.isAnyPropertySet()) {
+            producePersistenceUnitDescriptorFromConfig(
+                    hibernateOrmConfig, DEFAULT_PERSISTENCE_UNIT_NAME, hibernateOrmConfig.defaultPersistenceUnit,
+                    modelClassesPerPersistencesUnits.getOrDefault(DEFAULT_PERSISTENCE_UNIT_NAME, Collections.emptySet()),
+                    jdbcDataSources, applicationArchivesBuildItem, launchMode,
+                    systemProperties, nativeImageResources, hotDeploymentWatchedFiles, persistenceUnitDescriptors);
+        }
+
+        for (Entry<String, HibernateOrmConfigPersistenceUnit> persistenceUnitEntry : hibernateOrmConfig.persistenceUnits
+                .entrySet()) {
+            producePersistenceUnitDescriptorFromConfig(
+                    hibernateOrmConfig, persistenceUnitEntry.getKey(), persistenceUnitEntry.getValue(),
+                    modelClassesPerPersistencesUnits.getOrDefault(persistenceUnitEntry.getKey(), Collections.emptySet()),
+                    jdbcDataSources, applicationArchivesBuildItem, launchMode,
+                    systemProperties, nativeImageResources, hotDeploymentWatchedFiles, persistenceUnitDescriptors);
+        }
     }
 
     private static void producePersistenceUnitDescriptorFromConfig(
             HibernateOrmConfig hibernateOrmConfig,
             String persistenceUnitName,
             HibernateOrmConfigPersistenceUnit persistenceUnitConfig,
+            Set<String> modelClasses,
             List<JdbcDataSourceBuildItem> jdbcDataSources,
             ApplicationArchivesBuildItem applicationArchivesBuildItem,
             LaunchMode launchMode,
             BuildProducer<SystemPropertyBuildItem> systemProperties,
             BuildProducer<NativeImageResourceBuildItem> nativeImageResources,
             BuildProducer<HotDeploymentWatchedFileBuildItem> hotDeploymentWatchedFiles,
             BuildProducer<PersistenceUnitDescriptorBuildItem> persistenceUnitDescriptors) {
-
         // Find the associated datasource
         JdbcDataSourceBuildItem jdbcDataSource;
         String dataSource;
@@ -645,6 +670,12 @@ private static void producePersistenceUnitDescriptorFromConfig(
                                     persistenceUnitConfig.datasource.get(), persistenceUnitName)));
             dataSource = persistenceUnitConfig.datasource.get();
         } else {
+            if (!isDefaultPersistenceUnit(persistenceUnitName)) {
+                // if it's not the default persistence unit, we mandate a datasource to prevent common errors
+                throw new ConfigurationException(
+                        String.format("Datasource must be defined for persistence unit '%s'.", persistenceUnitName));
+            }
+
             jdbcDataSource = jdbcDataSources.stream()
                     .filter(i -> i.isDefault())
                     .findFirst()
@@ -666,6 +697,14 @@ private static void producePersistenceUnitDescriptorFromConfig(
         // we found one
         ParsedPersistenceXmlDescriptor descriptor = new ParsedPersistenceXmlDescriptor(null); //todo URL
         descriptor.setName(persistenceUnitName);
+
+        descriptor.setExcludeUnlistedClasses(true);
+        if (modelClasses.isEmpty()) {
+            LOG.warnf("Could not find any entities affected to the persistence unit '%s'.", persistenceUnitName);
+        } else {
+            descriptor.addClasses(new ArrayList<>(modelClasses));
+        }
+
         descriptor.setTransactionType(PersistenceUnitTransactionType.JTA);
         descriptor.getProperties().setProperty(AvailableSettings.DIALECT, dialect.get());
 
@@ -794,7 +833,8 @@ private static void producePersistenceUnitDescriptorFromConfig(
         }
 
         persistenceUnitDescriptors.produce(
-                new PersistenceUnitDescriptorBuildItem(descriptor, dataSource, getMultiTenancyStrategy(hibernateOrmConfig),
+                new PersistenceUnitDescriptorBuildItem(descriptor, dataSource,
+                        getMultiTenancyStrategy(hibernateOrmConfig),
                         false));
     }
 
@@ -871,4 +911,65 @@ private void enhanceEntities(final JpaEntitiesBuildItem domainObjects,
         }
     }
 
+    private Map<String, Set<String>> getModelClassesPerPersistenceUnits(HibernateOrmConfig hibernateOrmConfig,
+            JpaEntitiesBuildItem jpaEntities) {
+        if (hibernateOrmConfig.persistenceUnits.isEmpty()) {
+            // no named persistence units, all the entities will be associated with the default one
+            // so we don't need to split them
+            return Collections.singletonMap(DEFAULT_PERSISTENCE_UNIT_NAME, jpaEntities.getAllModelClassNames());
+        }
+
+        Map<String, Set<String>> modelClassesPerPersistenceUnits = new HashMap<>();
+        for (String modelClassName : jpaEntities.getAllModelClassNames()) {
+            String selectedPersistenceUnit = null;
+            int weight = -1;
+
+            // handle the default persistence unit
+            if (hibernateOrmConfig.defaultPersistenceUnit.packages.isPresent()) {
+                for (String pakkage : hibernateOrmConfig.defaultPersistenceUnit.packages.get()) {
+                    if (!pakkage.endsWith(".")) {
+                        pakkage = pakkage + ".";
+                    }
+                    if (modelClassName.startsWith(pakkage) && pakkage.length() > weight) {
+                        selectedPersistenceUnit = DEFAULT_PERSISTENCE_UNIT_NAME;
+                        weight = pakkage.length();
+                    }
+                }
+            } else {
+                // it will be a catch all
+                selectedPersistenceUnit = DEFAULT_PERSISTENCE_UNIT_NAME;
+                weight = 0;
+            }
+
+            for (Entry<String, HibernateOrmConfigPersistenceUnit> candidatePersistenceUnitEntry : hibernateOrmConfig.persistenceUnits
+                    .entrySet()) {
+                String candidatePersistenceUnitName = candidatePersistenceUnitEntry.getKey();
+                Set<String> candidatePersistenceUnitPackages = candidatePersistenceUnitEntry.getValue().packages
+                        .orElseThrow(() -> new ConfigurationException(String.format(
+                                "Packages must be configured for persistence unit '%s'.", candidatePersistenceUnitName)));
+
+                for (String pakkage : candidatePersistenceUnitPackages) {
+                    if (!pakkage.endsWith(".")) {
+                        pakkage = pakkage + ".";
+                    }
+                    if (modelClassName.startsWith(pakkage) && pakkage.length() > weight) {
+                        selectedPersistenceUnit = candidatePersistenceUnitName;
+                        weight = pakkage.length();
+                    }
+                }
+            }
+            if (selectedPersistenceUnit != null) {
+                modelClassesPerPersistenceUnits.putIfAbsent(selectedPersistenceUnit, new HashSet<>());
+                modelClassesPerPersistenceUnits.get(selectedPersistenceUnit).add(modelClassName);
+            } else {
+                LOG.warnf("Could not find a suitable persistence unit for model class '%s'.", modelClassName);
+            }
+        }
+
+        return modelClassesPerPersistenceUnits;
+    }
+
+    private static boolean isDefaultPersistenceUnit(String name) {
+        return DEFAULT_PERSISTENCE_UNIT_NAME.equals(name);
+    }
 }

extensions/hibernate-orm/deployment/src/test/java/io/quarkus/hibernate/orm/ExcludePersistenceXmlConfigTest.java

@@ -11,6 +11,7 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 import io.quarkus.arc.Arc;
+import io.quarkus.hibernate.orm.deployment.HibernateOrmProcessor;
 import io.quarkus.hibernate.orm.enhancer.Address;
 import io.quarkus.test.QuarkusUnitTest;
 
@@ -43,7 +44,7 @@ public void testPersistenceAndConfigTest() {
         Arc.container().requestContext().activate();
         try {
             // it is the default entity manager from application.properties, not templatePU from the persistence.xml
-            Assertions.assertEquals("default",
+            Assertions.assertEquals(HibernateOrmProcessor.DEFAULT_PERSISTENCE_UNIT_NAME,
                     entityManager.getEntityManagerFactory().getProperties().get("hibernate.ejb.persistenceUnitName"));
         } finally {
             Arc.container().requestContext().deactivate();

extensions/hibernate-orm/deployment/src/test/java/io/quarkus/hibernate/orm/HibernateMetricsTestCase.java

@@ -19,6 +19,7 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 import io.quarkus.arc.Arc;
+import io.quarkus.hibernate.orm.deployment.HibernateOrmProcessor;
 import io.quarkus.test.QuarkusUnitTest;
 
 /**
@@ -57,9 +58,9 @@ public void setNumber(Long number) {
     @Transactional
     public void testMetrics() {
         assertEquals(0L, getCounterValueOrNull("hibernate-orm.queries.executed",
-                new Tag("entityManagerFactory", "default")));
+                new Tag("entityManagerFactory", HibernateOrmProcessor.DEFAULT_PERSISTENCE_UNIT_NAME)));
         assertEquals(0L, getCounterValueOrNull("hibernate-orm.entities.inserted",
-                new Tag("entityManagerFactory", "default")));
+                new Tag("entityManagerFactory", HibernateOrmProcessor.DEFAULT_PERSISTENCE_UNIT_NAME)));
         Arc.container().requestContext().activate();
         try {
             DummyEntity entity = new DummyEntity();
@@ -68,9 +69,9 @@ public void testMetrics() {
             em.flush();
             em.createQuery("from DummyEntity e").getResultList();
             assertEquals(1L, getCounterValueOrNull("hibernate-orm.queries.executed",
-                    new Tag("entityManagerFactory", "default")));
+                    new Tag("entityManagerFactory", HibernateOrmProcessor.DEFAULT_PERSISTENCE_UNIT_NAME)));
             assertEquals(1L, getCounterValueOrNull("hibernate-orm.entities.inserted",
-                    new Tag("entityManagerFactory", "default")));
+                    new Tag("entityManagerFactory", HibernateOrmProcessor.DEFAULT_PERSISTENCE_UNIT_NAME)));
         } finally {
             Arc.container().requestContext().terminate();
         }