Allow connecting a Panache entity to more than one PU (at most one blocking, one reactive)

Author: FroMage

extensions/panache/hibernate-panache/deployment/src/main/java/io/quarkus/hibernate/panache/deployment/EntityToPersistenceUnitBuildItem.java

@@ -10,10 +10,13 @@ public final class EntityToPersistenceUnitBuildItem extends MultiBuildItem {
 
     private final String entityClass;
     private final String persistenceUnitName;
+    private final String reactivePersistenceUnitName;
 
-    public EntityToPersistenceUnitBuildItem(String entityClass, String persistenceUnitName) {
+    public EntityToPersistenceUnitBuildItem(String entityClass, String persistenceUnitName,
+            String reactivePersistenceUnitName) {
         this.entityClass = entityClass;
         this.persistenceUnitName = persistenceUnitName;
+        this.reactivePersistenceUnitName = reactivePersistenceUnitName;
     }
 
     public String getEntityClass() {
@@ -23,4 +26,8 @@ public String getEntityClass() {
     public String getPersistenceUnitName() {
         return persistenceUnitName;
     }
+
+    public String getReactivePersistenceUnitName() {
+        return reactivePersistenceUnitName;
+    }
 }

extensions/panache/hibernate-panache/deployment/src/main/java/io/quarkus/hibernate/panache/deployment/EntityToPersistenceUnitUtil.java

@@ -1,34 +1,47 @@
 package io.quarkus.hibernate.panache.deployment;
 
-import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.TreeMap;
 
 import io.quarkus.hibernate.orm.deployment.JpaModelPersistenceUnitMappingBuildItem;
+import io.quarkus.hibernate.orm.deployment.PersistenceUnitDescriptorBuildItem;
 
 //FIXME: duplicate with ORM and probably HR
 public final class EntityToPersistenceUnitUtil {
 
+    @FunctionalInterface
+    public interface TriConsumer {
+        void consume(String entity, String persistenceUnit, String reactivePersistenceUnit);
+    }
+
     private EntityToPersistenceUnitUtil() {
     }
 
     /**
      * Given the candidate entities, return a map with the persistence unit that single contains them
      * or throw an exception if any of the candidates is part of more than one persistence unit
      */
-    public static Map<String, String> determineEntityPersistenceUnits(
+    public static void determineEntityPersistenceUnits(
             Optional<JpaModelPersistenceUnitMappingBuildItem> jpaModelPersistenceUnitMapping,
-            Set<String> candidates, String source) {
+            List<PersistenceUnitDescriptorBuildItem> descriptors,
+            Set<String> candidates, String source,
+            TriConsumer consumer) {
         if (jpaModelPersistenceUnitMapping.isEmpty()) {
-            return Collections.emptyMap();
+            return;
         }
         Map<String, String> result = new HashMap<>();
         Map<String, Set<String>> collectedEntityToPersistenceUnits = jpaModelPersistenceUnitMapping.get()
                 .getEntityToPersistenceUnits();
 
+        Map<String, PersistenceUnitDescriptorBuildItem> descriptorsMap = new HashMap<>();
+        for (PersistenceUnitDescriptorBuildItem descriptor : descriptors) {
+            descriptorsMap.put(descriptor.getPersistenceUnitName(), descriptor);
+        }
+
         Map<String, Set<String>> violatingEntities = new TreeMap<>();
 
         for (Map.Entry<String, Set<String>> entry : collectedEntityToPersistenceUnits.entrySet()) {
@@ -40,10 +53,39 @@ public static Map<String, String> determineEntityPersistenceUnits(
                 continue;
             }
 
-            if (selectedPersistenceUnits.size() == 1) {
-                result.put(entityName, selectedPersistenceUnits.iterator().next());
-            } else {
-                violatingEntities.put(entityName, selectedPersistenceUnits);
+            String persistenceUnit = null;
+            String reactivePersistenceUnit = null;
+            if (selectedPersistenceUnits.size() > 0) {
+                boolean error = false;
+                // collect at most one of blocking/reactive PU
+                for (String selectedPersistenceUnit : selectedPersistenceUnits) {
+                    PersistenceUnitDescriptorBuildItem descriptor = descriptorsMap.get(selectedPersistenceUnit);
+                    if (descriptor == null) {
+                        throw new IllegalStateException(String.format("Cannot find descriptor unit named %s for entity %s",
+                                selectedPersistenceUnit, entityName));
+                    }
+                    if (descriptor.isReactive()) {
+                        if (reactivePersistenceUnit != null) {
+                            error = true;
+                            break;
+                        } else {
+                            reactivePersistenceUnit = selectedPersistenceUnit;
+                        }
+                    } else {
+                        if (persistenceUnit != null) {
+                            error = true;
+                            break;
+                        } else {
+                            persistenceUnit = selectedPersistenceUnit;
+                        }
+                    }
+                }
+                if (error) {
+                    violatingEntities.put(entityName, selectedPersistenceUnits);
+                } else if (persistenceUnit != null || reactivePersistenceUnit != null) {
+                    // ignore the entity if it belongs to no PU, though I doubt that happens
+                    consumer.consume(entityName, persistenceUnit, reactivePersistenceUnit);
+                }
             }
         }
 
@@ -57,7 +99,5 @@ public static Map<String, String> determineEntityPersistenceUnits(
                 throw new IllegalStateException(message.toString());
             }
         }
-
-        return result;
     }
 }

extensions/panache/hibernate-panache/deployment/src/main/java/io/quarkus/hibernate/panache/deployment/PanacheHibernateResourceProcessor.java

@@ -36,6 +36,7 @@
 import io.quarkus.deployment.builditem.FeatureBuildItem;
 import io.quarkus.deployment.util.JandexUtil;
 import io.quarkus.hibernate.orm.deployment.JpaModelPersistenceUnitMappingBuildItem;
+import io.quarkus.hibernate.orm.deployment.PersistenceUnitDescriptorBuildItem;
 import io.quarkus.hibernate.orm.deployment.spi.AdditionalJpaModelBuildItem;
 import io.quarkus.hibernate.panache.PanacheEntityMarker;
 import io.quarkus.hibernate.panache.PanacheRepositoryQueries;
@@ -112,6 +113,7 @@ void build(
             BuildProducer<BytecodeTransformerBuildItem> transformers,
             List<PanacheEntityClassBuildItem> entityClasses,
             Optional<JpaModelPersistenceUnitMappingBuildItem> jpaModelPersistenceUnitMapping,
+            List<PersistenceUnitDescriptorBuildItem> descriptors,
             List<PanacheMethodCustomizerBuildItem> methodCustomizersBuildItems,
             BuildProducer<EntityToPersistenceUnitBuildItem> entityToPersistenceUnit) {
 
@@ -136,8 +138,8 @@ void build(
 
         panacheEntities.addAll(modelClasses);
 
-        determineEntityPersistenceUnits(jpaModelPersistenceUnitMapping, panacheEntities, "Panache")
-                .forEach((e, pu) -> entityToPersistenceUnit.produce(new EntityToPersistenceUnitBuildItem(e, pu)));
+        determineEntityPersistenceUnits(jpaModelPersistenceUnitMapping, descriptors, panacheEntities, "Panache", (e, pu,
+                reactivePU) -> entityToPersistenceUnit.produce(new EntityToPersistenceUnitBuildItem(e, pu, reactivePU)));
     }
 
     @BuildStep
@@ -147,6 +149,7 @@ void recordEntityToPersistenceUnit(Optional<JpaModelPersistenceUnitMappingBuildI
             CombinedIndexBuildItem index,
             PanacheHibernateRecorder recorder) throws ClassNotFoundException {
         // PU
+        // FIXME: for now, this ignores the reactive PUs, but reactive PUs are not supported yet by Panache Reactive
         Map<String, String> map = new HashMap<>();
         for (EntityToPersistenceUnitBuildItem item : items) {
             map.put(item.getEntityClass(), item.getPersistenceUnitName());