Hibernate Reactive Panache: fix WithSessionOnDemand implementation

- fixes quarkusio#35568
mkouba · Aug 30, 2023 · 0d1ae5e · 0d1ae5e
1 parent 1347a17
commit 0d1ae5e
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 3 deletions.
diff --git a/...src/main/java/io/quarkus/hibernate/reactive/panache/common/runtime/SessionOperations.java b/...src/main/java/io/quarkus/hibernate/reactive/panache/common/runtime/SessionOperations.java
@@ -49,6 +49,7 @@ public Key<Session> get() {
 
     // This key is used to indicate that a reactive session should be opened lazily (when needed) in the current vertx context
     private static final String SESSION_ON_DEMAND_KEY = "hibernate.reactive.panache.sessionOnDemand";
+    private static final String SESSION_ON_DEMAND_OPENED_KEY = "hibernate.reactive.panache.sessionOnDemandOpened";
 
     /**
      * Marks the current vertx duplicated context as "lazy" which indicates that a reactive session should be opened lazily if
@@ -70,6 +71,7 @@ static <T> Uni<T> withSessionOnDemand(Supplier<Uni<T>> work) {
             // perform the work and eventually close the session and remove the key
             return work.get().eventually(() -> {
                 context.removeLocal(SESSION_ON_DEMAND_KEY);
+                context.removeLocal(SESSION_ON_DEMAND_OPENED_KEY);
                 return closeSession();
             });
         }
@@ -148,9 +150,15 @@ public static Uni<Mutiny.Session> getSession() {
             return Uni.createFrom().item(current);
         } else {
             if (context.getLocal(SESSION_ON_DEMAND_KEY) != null) {
-                // open a new reactive session and store it in the vertx duplicated context
-                // the context was marked as "lazy" which means that the session will be eventually closed
-                return getSessionFactory().openSession().invoke(s -> context.putLocal(key, s));
+                if (context.getLocal(SESSION_ON_DEMAND_OPENED_KEY) != null) {
+                    // a new reactive session is opened in a previous stage
+                    return Uni.createFrom().item(() -> getCurrentSession());
+                } else {
+                    // open a new reactive session and store it in the vertx duplicated context
+                    // the context was marked as "lazy" which means that the session will be eventually closed
+                    context.putLocal(SESSION_ON_DEMAND_OPENED_KEY, true);
+                    return getSessionFactory().openSession().invoke(s -> context.putLocal(key, s));
+                }
             } else {
                 throw new IllegalStateException("No current Mutiny.Session found"
                         + "\n\t- no reactive session was found in the context and the context was not marked to open a new session lazily"

diff --git a/extensions/panache/hibernate-reactive-panache/deployment/pom.xml b/extensions/panache/hibernate-reactive-panache/deployment/pom.xml
@@ -35,6 +35,11 @@
             <artifactId>quarkus-junit5-internal</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-test-hibernate-reactive-panache</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-reactive-pg-client-deployment</artifactId>

diff --git a/...ent/src/test/java/io/quarkus/hibernate/reactive/panache/test/WithSessionOnDemandTest.java b/...ent/src/test/java/io/quarkus/hibernate/reactive/panache/test/WithSessionOnDemandTest.java
@@ -0,0 +1,46 @@
+package io.quarkus.hibernate.reactive.panache.test;
+
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+
+import org.hibernate.reactive.mutiny.Mutiny.Session;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.arc.Unremovable;
+import io.quarkus.hibernate.reactive.panache.Panache;
+import io.quarkus.hibernate.reactive.panache.common.WithSessionOnDemand;
+import io.quarkus.test.QuarkusUnitTest;
+import io.quarkus.test.vertx.RunOnVertxContext;
+import io.quarkus.test.vertx.UniAsserter;
+import io.smallrye.mutiny.Uni;
+
+public class WithSessionOnDemandTest {
+
+    @RegisterExtension
+    static QuarkusUnitTest config = new QuarkusUnitTest()
+            .withApplicationRoot(root -> root.addClasses(MrBean.class, MyEntity.class));
+
+    @Inject
+    MrBean bean;
+
+    @RunOnVertxContext
+    @Test
+    public void testSharedSession(UniAsserter asserter) {
+        asserter.assertEquals(() -> bean.ping(), "true");
+    }
+
+    @Unremovable
+    @ApplicationScoped
+    static class MrBean {
+
+        @WithSessionOnDemand
+        Uni<String> ping() {
+            Uni<Session> s1 = Panache.getSession();
+            Uni<Session> s2 = Panache.getSession();
+            // Test that both unis receive the same session
+            return s1.chain(s1Val -> s2.map(s2Val -> "" + (s1Val == s2Val)));
+        }
+
+    }
+}