Panache Support multiple persistence unit in Hibernate Reactive

* Integration test for multiple reactive persistence units and Panache
Backported from ORM the handling of different persistence units in entity in Panache
* `@WithSessionOnDemand` works only with the default persistence unit
* withSession overload to take the PU name
* Execute update without entity runs on default session
* Repository flush() flushes all PUs
Refactor
* Use ComputingCache instead of a Map<String, LazyValue>
* Keep track of onDemand created sessions to close them accordingly
* create map of entity => PU in kotlin as well (identical to what happens in Hibernate ORM)
* Aligned reactive blocking to orm repositories by using @GenerateBridge to call getSession
* getSession shouldn't be static to be substituted by @GenerateBridge
* Added Panache.withTransaction overload to select persistence unit
* Updated Panache documentation
* Support @WithSession("pu-name") and @WithTransaction("pu-name")
* Added test to verify the same session is used
* Removed dependency from Kotlin panache module to Java, created two new Panache classes inside the Kotlin module and changed the tests accordingly.

Author: lucamolteni

docs/src/main/asciidoc/hibernate-reactive-panache.adoc

@@ -772,7 +772,10 @@ If no annotations are found.
 
 == Multiple Persistence Units
 
-Hibernate Reactive in Quarkus currently does not support multiple persistence units.
+Hibernate Reactive with Panache supports multiple persistence units. 
+When using multiple persistence units, Panache automatically uses the correct persistence unit based on the entity's configuration.
+
+For configuration details on setting up multiple persistence units, refer to the xref:hibernate-orm.adoc#multiple-persistence-units[Multiple Persistence Units] section in the Hibernate ORM guide.
 
 [[transactions]]
 == Sessions and Transactions
@@ -783,13 +786,15 @@ For example, if a Panache entity method is invoked in a Jakarta REST resource me
 For other cases, there are both a declarative and a programmatic way to ensure the session is opened.
 You can annotate a CDI business method that returns `Uni` with the `@WithSession` annotation.
 The method will be intercepted and the returned `Uni` will be triggered within a scope of a reactive session.
+If you have multiple persistence units, you can specify which one to use by providing the persistence unit name as the annotation value: `@WithSession("my-persistence-unit")`. If not specified, the default persistence unit is used.
 Alternatively, you can use the `Panache.withSession()` method to achieve the same effect.
 
 NOTE: Note that a Panache entity may not be used from a blocking thread. See also xref:getting-started-reactive.adoc[Getting Started With Reactive] guide that explains the basics of reactive principles in Quarkus.
 
 Also make sure to wrap methods that modify the database or involve multiple queries (e.g. `entity.persist()`) within a transaction.
 You can annotate a CDI business method that returns `Uni` with the `@WithTransaction` annotation.
 The method will be intercepted and the returned `Uni` is triggered within a transaction boundary.
+If you have multiple persistence units, you can specify which one to use by providing the persistence unit name as the annotation value: `@WithTransaction("my-persistence-unit")`. If not specified, the default persistence unit is used.
 Alternatively, you can use the `Panache.withTransaction()` method for the same effect.
 
 IMPORTANT: You cannot use the `@Transactional` annotation with Hibernate Reactive for your transactions: you must use `@WithTransaction`, and your annotated method must return a `Uni` to be non-blocking.

extensions/panache/hibernate-reactive-panache-common/runtime/src/main/java/io/quarkus/hibernate/reactive/panache/common/WithSession.java

@@ -1,11 +1,14 @@
 package io.quarkus.hibernate.reactive.panache.common;
 
+import static io.quarkus.hibernate.orm.runtime.PersistenceUnitUtil.DEFAULT_PERSISTENCE_UNIT_NAME;
+
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Inherited;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import jakarta.enterprise.util.Nonbinding;
 import jakarta.interceptor.InterceptorBinding;
 
 /**
@@ -25,4 +28,12 @@
 @Retention(value = RetentionPolicy.RUNTIME)
 public @interface WithSession {
 
+    /**
+     * The name of the persistence unit. If not specified, the default persistence unit is used.
+     *
+     * @return the persistence unit name
+     */
+    @Nonbinding
+    String value() default DEFAULT_PERSISTENCE_UNIT_NAME;
+
 }

extensions/panache/hibernate-reactive-panache-common/runtime/src/main/java/io/quarkus/hibernate/reactive/panache/common/WithTransaction.java

@@ -1,11 +1,14 @@
 package io.quarkus.hibernate.reactive.panache.common;
 
+import static io.quarkus.hibernate.orm.runtime.PersistenceUnitUtil.DEFAULT_PERSISTENCE_UNIT_NAME;
+
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Inherited;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import jakarta.enterprise.util.Nonbinding;
 import jakarta.interceptor.InterceptorBinding;
 
 /**
@@ -27,4 +30,12 @@
 @Retention(value = RetentionPolicy.RUNTIME)
 public @interface WithTransaction {
 
+    /**
+     * The name of the persistence unit. If not specified, the default persistence unit is used.
+     *
+     * @return the persistence unit name
+     */
+    @Nonbinding
+    String value() default DEFAULT_PERSISTENCE_UNIT_NAME;
+
 }

extensions/panache/hibernate-reactive-panache-common/runtime/src/main/java/io/quarkus/hibernate/reactive/panache/common/runtime/AbstractJpaOperations.java

@@ -1,10 +1,14 @@
 package io.quarkus.hibernate.reactive.panache.common.runtime;
 
+import static io.quarkus.hibernate.orm.runtime.PersistenceUnitUtil.DEFAULT_PERSISTENCE_UNIT_NAME;
+
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import jakarta.persistence.LockModeType;
@@ -19,6 +23,11 @@
 import io.smallrye.mutiny.Uni;
 
 public abstract class AbstractJpaOperations<PanacheQueryType> {
+    private static volatile Map<String, String> entityToPersistenceUnit = Collections.emptyMap();
+
+    public static void setEntityToPersistenceUnit(Map<String, String> map) {
+        entityToPersistenceUnit = Collections.unmodifiableMap(map);
+    }
 
     // FIXME: make it configurable?
     static final long TIMEOUT_MS = 5000;
@@ -34,7 +43,7 @@ protected abstract PanacheQueryType createPanacheQuery(Uni<Mutiny.Session> sessi
     // Instance methods
 
     public Uni<Void> persist(Object entity) {
-        return persist(getSession(), entity);
+        return persist(getSession(entity.getClass()), entity);
     }
 
     public Uni<Void> persist(Uni<Mutiny.Session> sessionUni, Object entity) {
@@ -67,20 +76,38 @@ public Uni<Void> persist(Stream<?> entities) {
     }
 
     public Uni<Void> persist(Object... entities) {
-        return getSession().chain(session -> session.persistAll(entities));
+        Map<String, List<Object>> sessions = Arrays.stream(entities)
+                .collect(Collectors.groupingBy(e -> entityToPersistenceUnit.get(e.getClass().getName())));
+
+        List<Uni<Void>> results = new ArrayList<>();
+        for (Entry<String, List<Object>> entry : sessions.entrySet()) {
+            results.add(getSession(entry.getKey()).chain(session -> session.persistAll(entry.getValue().toArray())));
+        }
+
+        return Uni.combine().all().unis(results).discardItems();
     }
 
     public Uni<Void> delete(Object entity) {
-        return getSession().chain(session -> session.remove(entity));
+        return getSession(entity.getClass()).chain(session -> session.remove(entity));
     }
 
     public boolean isPersistent(Object entity) {
-        Mutiny.Session current = SessionOperations.getCurrentSession();
-        return current != null ? current.contains(entity) : false;
+        Session currentSession = getCurrentSession(entity.getClass());
+        if (currentSession == null) {
+            // No active session so object is surely non-persistent
+            return false;
+        }
+
+        return currentSession.contains(entity);
+    }
+
+    public Session getCurrentSession(Class<?> entityClass) {
+        String persistenceUnitName = entityToPersistenceUnit.get(entityClass.getName());
+        return SessionOperations.getCurrentSession(persistenceUnitName);
     }
 
-    public Uni<Void> flush() {
-        return getSession().chain(Session::flush);
+    public Uni<Void> flush(Object entity) {
+        return getSession(entity.getClass()).chain(Session::flush);
     }
 
     public int paramCount(Object[] params) {
@@ -95,11 +122,11 @@ public int paramCount(Map<String, Object> params) {
     // Queries
 
     public Uni<?> findById(Class<?> entityClass, Object id) {
-        return getSession().chain(session -> session.find(entityClass, id));
+        return getSession(entityClass).chain(session -> session.find(entityClass, id));
     }
 
     public Uni<?> findById(Class<?> entityClass, Object id, LockModeType lockModeType) {
-        return getSession()
+        return getSession(entityClass)
                 .chain(session -> session.find(entityClass, id, LockModeConverter.convertToLockMode(lockModeType)));
     }
 
@@ -108,7 +135,7 @@ public PanacheQueryType find(Class<?> entityClass, String panacheQuery, Object..
     }
 
     public PanacheQueryType find(Class<?> entityClass, String panacheQuery, Sort sort, Object... params) {
-        Uni<Mutiny.Session> session = getSession();
+        Uni<Mutiny.Session> session = getSession(entityClass);
         if (PanacheJpaUtil.isNamedQuery(panacheQuery)) {
             String namedQuery = panacheQuery.substring(1);
             if (sort != null) {
@@ -128,7 +155,7 @@ public PanacheQueryType find(Class<?> entityClass, String panacheQuery, Map<Stri
     }
 
     public PanacheQueryType find(Class<?> entityClass, String panacheQuery, Sort sort, Map<String, Object> params) {
-        Uni<Mutiny.Session> session = getSession();
+        Uni<Mutiny.Session> session = getSession(entityClass);
         if (PanacheJpaUtil.isNamedQuery(panacheQuery)) {
             String namedQuery = panacheQuery.substring(1);
             if (sort != null) {
@@ -177,13 +204,13 @@ public Uni<List<?>> list(Class<?> entityClass, String query, Sort sort, Paramete
 
     public PanacheQueryType findAll(Class<?> entityClass) {
         String query = "FROM " + PanacheJpaUtil.getEntityName(entityClass);
-        Uni<Mutiny.Session> session = getSession();
+        Uni<Mutiny.Session> session = getSession(entityClass);
         return createPanacheQuery(session, query, null, null, null);
     }
 
     public PanacheQueryType findAll(Class<?> entityClass, Sort sort) {
         String query = "FROM " + PanacheJpaUtil.getEntityName(entityClass);
-        Uni<Mutiny.Session> session = getSession();
+        Uni<Mutiny.Session> session = getSession(entityClass);
         return createPanacheQuery(session, query, null, PanacheJpaUtil.toOrderBy(sort), null);
     }
 
@@ -196,7 +223,7 @@ public Uni<List<?>> listAll(Class<?> entityClass, Sort sort) {
     }
 
     public Uni<Long> count(Class<?> entityClass) {
-        return getSession()
+        return getSession(entityClass)
                 .chain(session -> session
                         .createSelectionQuery("FROM " + PanacheJpaUtil.getEntityName(entityClass), entityClass)
                         .getResultCount());
@@ -206,13 +233,13 @@ public Uni<Long> count(Class<?> entityClass) {
     public Uni<Long> count(Class<?> entityClass, String panacheQuery, Object... params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return (Uni) getSession().chain(session -> {
+            return (Uni) getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName, Long.class), params).getSingleResult();
             });
 
-        return getSession().chain(session -> bindParameters(
+        return getSession(entityClass).chain(session -> bindParameters(
                 session.createSelectionQuery(PanacheJpaUtil.createQueryForCount(entityClass, panacheQuery, paramCount(params)),
                         Object.class),
                 params).getResultCount())
@@ -223,13 +250,13 @@ public Uni<Long> count(Class<?> entityClass, String panacheQuery, Object... para
     public Uni<Long> count(Class<?> entityClass, String panacheQuery, Map<String, Object> params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return getSession().chain(session -> {
+            return getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName, Long.class), params).getSingleResult();
             });
 
-        return getSession().chain(session -> bindParameters(
+        return getSession(entityClass).chain(session -> bindParameters(
                 session.createSelectionQuery(PanacheJpaUtil.createQueryForCount(entityClass, panacheQuery, paramCount(params)),
                         Object.class),
                 params).getResultCount())
@@ -258,7 +285,7 @@ public Uni<Boolean> exists(Class<?> entityClass, String query, Parameters params
     }
 
     public Uni<Long> deleteAll(Class<?> entityClass) {
-        return getSession().chain(
+        return getSession(entityClass).chain(
                 session -> session.createMutationQuery("DELETE FROM " + PanacheJpaUtil.getEntityName(entityClass))
                         .executeUpdate()
                         .map(Integer::longValue));
@@ -272,20 +299,20 @@ public Uni<Boolean> deleteById(Class<?> entityClass, Object id) {
                     if (entity == null) {
                         return Uni.createFrom().item(false);
                     }
-                    return getSession().chain(session -> session.remove(entity).map(v -> true));
+                    return getSession(entityClass).chain(session -> session.remove(entity).map(v -> true));
                 });
     }
 
     public Uni<Long> delete(Class<?> entityClass, String panacheQuery, Object... params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return getSession().chain(session -> {
+            return getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName), params).executeUpdate().map(Integer::longValue);
             });
 
-        return getSession().chain(session -> bindParameters(
+        return getSession(entityClass).chain(session -> bindParameters(
                 session.createMutationQuery(PanacheJpaUtil.createDeleteQuery(entityClass, panacheQuery, paramCount(params))),
                 params)
                 .executeUpdate().map(Integer::longValue))
@@ -296,13 +323,13 @@ public Uni<Long> delete(Class<?> entityClass, String panacheQuery, Object... par
     public Uni<Long> delete(Class<?> entityClass, String panacheQuery, Map<String, Object> params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return getSession().chain(session -> {
+            return getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName), params).executeUpdate().map(Integer::longValue);
             });
 
-        return getSession().chain(session -> bindParameters(
+        return getSession(entityClass).chain(session -> bindParameters(
                 session.createMutationQuery(PanacheJpaUtil.createDeleteQuery(entityClass, panacheQuery, paramCount(params))),
                 params)
                 .executeUpdate().map(Integer::longValue))
@@ -314,15 +341,15 @@ public Uni<Long> delete(Class<?> entityClass, String query, Parameters params) {
         return delete(entityClass, query, params.map());
     }
 
-    public IllegalStateException implementationInjectionMissing() {
+    public static IllegalStateException implementationInjectionMissing() {
         return new IllegalStateException(
                 "This method is normally automatically overridden in subclasses: did you forget to annotate your entity with @Entity?");
     }
 
     public Uni<Integer> executeUpdate(Class<?> entityClass, String panacheQuery, Object... params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return (Uni) getSession().chain(session -> {
+            return (Uni) getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName), params).executeUpdate();
@@ -337,7 +364,7 @@ public Uni<Integer> executeUpdate(Class<?> entityClass, String panacheQuery, Obj
     public Uni<Integer> executeUpdate(Class<?> entityClass, String panacheQuery, Map<String, Object> params) {
 
         if (PanacheJpaUtil.isNamedQuery(panacheQuery))
-            return (Uni) getSession().chain(session -> {
+            return (Uni) getSession(entityClass).chain(session -> {
                 String namedQueryName = panacheQuery.substring(1);
                 NamedQueryUtil.checkNamedQuery(entityClass, namedQueryName);
                 return bindParameters(session.createNamedQuery(namedQueryName), params).executeUpdate();
@@ -364,8 +391,18 @@ public Uni<Integer> update(Class<?> entityClass, String query, Object... params)
     //
     // Static helpers
 
-    public static Uni<Mutiny.Session> getSession() {
-        return SessionOperations.getSession();
+    public Uni<Mutiny.Session> getSession() {
+        return getSession(DEFAULT_PERSISTENCE_UNIT_NAME);
+    }
+
+    public Uni<Mutiny.Session> getSession(Class<?> clazz) {
+        String className = clazz.getName();
+        String persistenceUnitName = entityToPersistenceUnit.get(className);
+        return getSession(persistenceUnitName);
+    }
+
+    public Uni<Mutiny.Session> getSession(String persistenceUnitName) {
+        return SessionOperations.getSession(persistenceUnitName);
     }
 
     public static Mutiny.Query<?> bindParameters(Mutiny.Query<?> query, Object[] params) {
@@ -395,13 +432,21 @@ public static <T extends Mutiny.AbstractQuery> T bindParameters(T query, Map<Str
         return query;
     }
 
-    public static Uni<Integer> executeUpdate(String query, Object... params) {
-        return getSession().chain(session -> bindParameters(session.createMutationQuery(query), params)
-                .executeUpdate());
+    /**
+     * Execute update on default persistence unit
+     */
+    public Uni<Integer> executeUpdate(String query, Object... params) {
+        return getSession(DEFAULT_PERSISTENCE_UNIT_NAME)
+                .chain(session -> bindParameters(session.createMutationQuery(query), params)
+                        .executeUpdate());
     }
 
-    public static Uni<Integer> executeUpdate(String query, Map<String, Object> params) {
-        return getSession().chain(session -> bindParameters(session.createMutationQuery(query), params)
-                .executeUpdate());
+    /**
+     * Execute update on default persistence unit
+     */
+    public Uni<Integer> executeUpdate(String query, Map<String, Object> params) {
+        return getSession(DEFAULT_PERSISTENCE_UNIT_NAME)
+                .chain(session -> bindParameters(session.createMutationQuery(query), params)
+                        .executeUpdate());
     }
 }