Integration test for multiple reactive persistence units and Panache

Backported from ORM the handling of different persistence units in entity in Panache

Author: lucamolteni

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

@@ -0,0 +1,25 @@
+package io.quarkus.hibernate.reactive.panache.deployment;
+
+import io.quarkus.builder.item.MultiBuildItem;
+
+/**
+ * Used to record that a specific JPA entity is associated with a specific persistence unit
+ */
+public final class EntityToPersistenceUnitBuildItem extends MultiBuildItem {
+
+    private final String entityClass;
+    private final String persistenceUnitName;
+
+    public EntityToPersistenceUnitBuildItem(String entityClass, String persistenceUnitName) {
+        this.entityClass = entityClass;
+        this.persistenceUnitName = persistenceUnitName;
+    }
+
+    public String getEntityClass() {
+        return entityClass;
+    }
+
+    public String getPersistenceUnitName() {
+        return persistenceUnitName;
+    }
+}

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

@@ -0,0 +1,62 @@
+package io.quarkus.hibernate.reactive.panache.deployment;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.TreeMap;
+
+import io.quarkus.hibernate.orm.deployment.JpaModelPersistenceUnitMappingBuildItem;
+
+public final class EntityToPersistenceUnitUtil {
+
+    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(
+            Optional<JpaModelPersistenceUnitMappingBuildItem> jpaModelPersistenceUnitMapping,
+            Set<String> candidates, String source) {
+        if (jpaModelPersistenceUnitMapping.isEmpty()) {
+            return Collections.emptyMap();
+        }
+        Map<String, String> result = new HashMap<>();
+        Map<String, Set<String>> collectedEntityToPersistenceUnits = jpaModelPersistenceUnitMapping.get()
+                .getEntityToPersistenceUnits();
+
+        Map<String, Set<String>> violatingEntities = new TreeMap<>();
+
+        for (Map.Entry<String, Set<String>> entry : collectedEntityToPersistenceUnits.entrySet()) {
+            String entityName = entry.getKey();
+            Set<String> selectedPersistenceUnits = entry.getValue();
+            boolean isCandidate = candidates.contains(entityName);
+
+            if (!isCandidate) {
+                continue;
+            }
+
+            if (selectedPersistenceUnits.size() == 1) {
+                result.put(entityName, selectedPersistenceUnits.iterator().next());
+            } else {
+                violatingEntities.put(entityName, selectedPersistenceUnits);
+            }
+        }
+
+        if (violatingEntities.size() > 0) {
+            StringBuilder message = new StringBuilder(
+                    String.format("%s entities do not support being attached to several persistence units:\n", source));
+            for (Map.Entry<String, Set<String>> violatingEntityEntry : violatingEntities
+                    .entrySet()) {
+                message.append("\t- ").append(violatingEntityEntry.getKey()).append(" is attached to: ")
+                        .append(String.join(",", violatingEntityEntry.getValue()));
+                throw new IllegalStateException(message.toString());
+            }
+        }
+
+        return result;
+    }
+}

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

@@ -1,8 +1,11 @@
 package io.quarkus.hibernate.reactive.panache.deployment;
 
+import static io.quarkus.hibernate.reactive.panache.deployment.EntityToPersistenceUnitUtil.determineEntityPersistenceUnits;
+
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -28,6 +31,7 @@
 import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
 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.spi.AdditionalJpaModelBuildItem;
 import io.quarkus.hibernate.reactive.panache.PanacheEntity;
 import io.quarkus.hibernate.reactive.panache.PanacheEntityBase;
@@ -94,19 +98,25 @@ void collectEntityClasses(CombinedIndexBuildItem index, BuildProducer<PanacheEnt
     void build(CombinedIndexBuildItem index,
             BuildProducer<BytecodeTransformerBuildItem> transformers,
             List<PanacheEntityClassBuildItem> entityClasses,
-            List<PanacheMethodCustomizerBuildItem> methodCustomizersBuildItems) throws Exception {
+            Optional<JpaModelPersistenceUnitMappingBuildItem> jpaModelPersistenceUnitMapping,
+            List<PanacheMethodCustomizerBuildItem> methodCustomizersBuildItems,
+            BuildProducer<EntityToPersistenceUnitBuildItem> entityToPersistenceUnit) throws Exception {
 
         List<PanacheMethodCustomizer> methodCustomizers = methodCustomizersBuildItems.stream()
                 .map(bi -> bi.getMethodCustomizer()).collect(Collectors.toList());
 
         PanacheJpaRepositoryEnhancer daoEnhancer = new PanacheJpaRepositoryEnhancer(index.getIndex());
         Set<String> daoClasses = new HashSet<>();
+        Set<String> panacheEntities = new HashSet<>();
         for (ClassInfo classInfo : index.getIndex().getAllKnownImplementors(DOTNAME_PANACHE_REPOSITORY_BASE)) {
             // Skip PanacheRepository
             if (classInfo.name().equals(DOTNAME_PANACHE_REPOSITORY))
                 continue;
             if (daoEnhancer.skipRepository(classInfo))
                 continue;
+            List<org.jboss.jandex.Type> typeParameters = JandexUtil
+                    .resolveTypeParameters(classInfo.name(), DOTNAME_PANACHE_REPOSITORY_BASE, index.getIndex());
+            panacheEntities.add(typeParameters.get(0).name().toString());
             daoClasses.add(classInfo.name().toString());
         }
         for (ClassInfo classInfo : index.getIndex().getAllKnownImplementors(DOTNAME_PANACHE_REPOSITORY)) {
@@ -121,10 +131,20 @@ void build(CombinedIndexBuildItem index,
         PanacheJpaEntityOperationsEnhancer entityOperationsEnhancer = new PanacheJpaEntityOperationsEnhancer(index.getIndex(),
                 methodCustomizers,
                 ReactiveJavaJpaTypeBundle.BUNDLE);
+
+        Set<String> modelClasses = new HashSet<>();
         for (PanacheEntityClassBuildItem entityClass : entityClasses) {
             String entityClassName = entityClass.get().name().toString();
+            modelClasses.add(entityClassName);
             transformers.produce(new BytecodeTransformerBuildItem(entityClassName, entityOperationsEnhancer));
         }
+
+        panacheEntities.addAll(modelClasses);
+
+        determineEntityPersistenceUnits(jpaModelPersistenceUnitMapping, panacheEntities, "Panache")
+                .forEach((e, pu) -> {
+                    entityToPersistenceUnit.produce(new EntityToPersistenceUnitBuildItem(e, pu));
+                });
     }
 
     @BuildStep

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/DefaultPersistenceUnitConfigTest.java

@@ -0,0 +1,34 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu;
+
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.first.FirstEntity;
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.second.SecondEntity;
+import io.quarkus.test.QuarkusUnitTest;
+import io.restassured.RestAssured;
+
+public class DefaultPersistenceUnitConfigTest {
+
+    @RegisterExtension
+    static QuarkusUnitTest runner = new QuarkusUnitTest()
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(FirstEntity.class, SecondEntity.class, PanacheTestResource.class)
+                    .addAsResource("application.properties", "application.properties"));
+
+    @Test
+    public void panacheOperations() {
+        /**
+         * First entity operations
+         */
+        RestAssured.when().get("/persistence-unit/first/name-1").then().body(Matchers.is("1"));
+        RestAssured.when().get("/persistence-unit/first/name-2").then().body(Matchers.is("2"));
+
+        /**
+         * second entity operations
+         */
+        RestAssured.when().get("/persistence-unit/second/name-1").then().body(Matchers.is("1"));
+        RestAssured.when().get("/persistence-unit/second/name-2").then().body(Matchers.is("2"));
+    }
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/ErroneousPersistenceUnitConfigTest.java

@@ -0,0 +1,24 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.first.FirstEntity;
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.second.SecondEntity;
+import io.quarkus.test.QuarkusUnitTest;
+
+public class ErroneousPersistenceUnitConfigTest {
+
+    @RegisterExtension
+    static QuarkusUnitTest runner = new QuarkusUnitTest()
+            .setExpectedException(IllegalStateException.class)
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(FirstEntity.class, SecondEntity.class, PanacheTestResource.class)
+                    .addAsResource("application-erroneous-multiple-persistence-units.properties", "application.properties"));
+
+    @Test
+    public void shouldNotReachHere() {
+        Assertions.fail();
+    }
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/MultiplePersistenceUnitConfigTest.java

@@ -0,0 +1,36 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu;
+
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.first.FirstEntity;
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.second.SecondEntity;
+import io.quarkus.test.QuarkusUnitTest;
+import io.restassured.RestAssured;
+
+public class MultiplePersistenceUnitConfigTest {
+
+    @RegisterExtension
+    static QuarkusUnitTest runner = new QuarkusUnitTest()
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(FirstEntity.class, SecondEntity.class, PanacheTestResource.class)
+                    .addAsResource("application-multiple-persistence-units.properties", "application.properties"));
+
+    @Test
+    public void panacheOperations() {
+        /**
+         * First entity operations
+         */
+        RestAssured.when().get("/persistence-unit/first/name-1").then().body(Matchers.is("1"));
+        RestAssured.when().get("/persistence-unit/first/name-2").then().body(Matchers.is("2"));
+
+        /**
+         * second entity operations
+         */
+        RestAssured.when().get("/persistence-unit/second/name-1").then().body(Matchers.is("1"));
+        RestAssured.when().get("/persistence-unit/second/name-2").then().body(Matchers.is("2"));
+    }
+
+    // Do a WithSessionOnDemandTest with multiple persistence
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/PanacheTestResource.java

@@ -0,0 +1,40 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu;
+
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.PathParam;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.core.MediaType;
+
+import io.quarkus.hibernate.reactive.panache.Panache;
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.first.FirstEntity;
+import io.quarkus.hibernate.reactive.panache.test.multiple_pu.second.SecondEntity;
+import io.smallrye.mutiny.Uni;
+
+@Path("/persistence-unit")
+public class PanacheTestResource {
+
+    @GET
+    @Path("/first/{name}")
+    @Produces(MediaType.TEXT_PLAIN)
+    public Uni<Long> createWithFirstPuAndReturnCount(@PathParam("name") String name) {
+        FirstEntity entity = new FirstEntity();
+        entity.name = name;
+
+        return Panache.withTransaction(() -> entity.persistAndFlush())
+                .flatMap(e -> FirstEntity.count());
+
+    }
+
+    @GET
+    @Path("/second/{name}")
+    @Produces(MediaType.TEXT_PLAIN)
+    public Uni<Long> createWithSecondPUAndReturnCount(@PathParam("name") String name) {
+        SecondEntity entity = new SecondEntity();
+        entity.name = name;
+
+        return Panache.withTransaction(() -> entity.persistAndFlush())
+                .flatMap(e -> SecondEntity.count());
+
+    }
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/first/FirstEntity.java

@@ -0,0 +1,11 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu.first;
+
+import jakarta.persistence.Entity;
+
+import io.quarkus.hibernate.reactive.panache.PanacheEntity;
+
+@Entity
+public class FirstEntity extends PanacheEntity {
+
+    public String name;
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/java/io/quarkus/hibernate/reactive/panache/test/multiple_pu/second/SecondEntity.java

@@ -0,0 +1,11 @@
+package io.quarkus.hibernate.reactive.panache.test.multiple_pu.second;
+
+import jakarta.persistence.Entity;
+
+import io.quarkus.hibernate.reactive.panache.PanacheEntity;
+
+@Entity
+public class SecondEntity extends PanacheEntity {
+
+    public String name;
+}

extensions/panache/hibernate-reactive-panache/deployment/src/test/resources/application-erroneous-multiple-persistence-units.properties

@@ -0,0 +1,18 @@
+quarkus.datasource.db-kind=postgresql
+quarkus.datasource.username=hibernate_orm_test
+quarkus.datasource.password=hibernate_orm_test
+quarkus.datasource.reactive.url=${postgres.reactive.url}
+quarkus.datasource.devservices.enabled=false
+
+quarkus.hibernate-orm.schema-management.strategy=drop-and-create
+quarkus.hibernate-orm.packages=io.quarkus.hibernate.reactive.panache.test.multiple_pu.first
+
+quarkus.datasource."second".db-kind=postgresql
+quarkus.datasource."second".username=hibernate_orm_test
+quarkus.datasource."second".password=hibernate_orm_test
+quarkus.datasource."second".reactive.url=${postgres.reactive.url}
+quarkus.datasource."second".devservices.enabled=false
+
+quarkus.hibernate-orm."second".datasource=second
+quarkus.hibernate-orm."second".schema-management.strategy=drop-and-create
+quarkus.hibernate-orm."second".packages=io.quarkus.hibernate.reactive.panache.test.multiple_pu.first