Hardening transaction state

hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySessionImpl.java

@@ -508,7 +508,8 @@ Uni<T> execute(Function<Mutiny.Transaction, Uni<T>> work) {
 		 * roll back) and an error thrown by the work.
 		 */
 		Uni<T> executeInTransaction(Function<Mutiny.Transaction, Uni<T>> work) {
-			return work.apply( this )
+			return Uni.createFrom()
+					.deferred( () -> work.apply( this ) )
 					// only flush() if the work completed with no exception
 					.call( this::flush ).call( this::beforeCompletion )
 					// in the case of an exception or cancellation

hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinyStatelessSessionImpl.java

@@ -267,7 +267,7 @@ Uni<T> execute(Function<Mutiny.Transaction, Uni<T>> work) {
 		 * and an error thrown by the work.
 		 */
 		Uni<T> executeInTransaction(Function<Mutiny.Transaction, Uni<T>> work) {
-			return work.apply( this )
+			return Uni.createFrom().deferred( () -> work.apply( this ) )
 					// in the case of an exception or cancellation
 					// we need to rollback the transaction
 					.onFailure().call( this::rollback )

hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientConnection.java

@@ -286,8 +286,12 @@ private SqlConnection client() {
 
 	@Override
 	public CompletionStage<Void> beginTransaction() {
+		if ( transaction != null ) {
+			throw new IllegalStateException( "Can't begin a new transaction as an active transaction is already associated to this connection" );
+		}
 		return connection.begin()
 				.onSuccess( tx -> LOG.tracef( "Transaction started: %s", tx ) )
+				.onFailure( v -> LOG.errorf( "Failed to start a transaction: %s", transaction ) )
 				.toCompletionStage()
 				.thenAccept( this::setTransaction );
 	}
@@ -296,22 +300,28 @@ public CompletionStage<Void> beginTransaction() {
 	public CompletionStage<Void> commitTransaction() {
 		return transaction.commit()
 				.onSuccess( v -> LOG.tracef( "Transaction committed: %s", transaction ) )
+				.onFailure( v -> LOG.errorf( "Failed to commit transaction: %s", transaction ) )
 				.toCompletionStage()
 				.whenComplete( this::clearTransaction );
 	}
 
 	@Override
 	public CompletionStage<Void> rollbackTransaction() {
 		return transaction.rollback()
+				.onFailure( v -> LOG.errorf( "Failed to rollback transaction: %s", transaction ) )
 				.onSuccess( v -> LOG.tracef( "Transaction rolled back: %s", transaction ) )
 				.toCompletionStage()
 				.whenComplete( this::clearTransaction );
 	}
 
 	@Override
 	public CompletionStage<Void> close() {
+		if ( transaction != null ) {
+			throw new IllegalStateException( "Connection being closed with a live transaction associated to it" );
+		}
 		return connection.close()
 				.onSuccess( event -> LOG.tracef( "Connection closed: %s", connection ) )
+				.onFailure( v -> LOG.errorf( "Failed to close a connection: %s", connection ) )
 				.toCompletionStage();
 	}
 
@@ -357,6 +367,7 @@ private void setTransaction(Transaction tx) {
 	}
 
 	private void clearTransaction(Void v, Throwable x) {
+		LOG.tracef( "Clearing current transaction instance from connection: %s", transaction );
 		transaction = null;
 	}
 

hibernate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageSessionImpl.java

@@ -46,6 +46,7 @@
 
 import static org.hibernate.reactive.util.impl.CompletionStages.applyToAll;
 import static org.hibernate.reactive.util.impl.CompletionStages.returnOrRethrow;
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
 /**
  * Implements the {@link Stage.Session} API. This delegating class is
@@ -434,7 +435,7 @@ CompletionStage<T> execute(Function<Stage.Transaction, CompletionStage<T>> work)
 		 * roll back) and an error thrown by the work.
 		 */
 		CompletionStage<T> executeInTransaction(Function<Stage.Transaction, CompletionStage<T>> work) {
-			return work.apply( this )
+			return voidFuture().thenCompose( v -> work.apply( this ) )
 					// only flush() if the work completed with no exception
 					.thenCompose( result -> flush().thenApply( v -> result ) )
 					// have to capture the error here and pass it along,

hibernate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageStatelessSessionImpl.java

@@ -26,6 +26,7 @@
 import java.util.function.Function;
 
 import static org.hibernate.reactive.util.impl.CompletionStages.returnOrRethrow;
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
 /**
  * Implements the {@link Stage.StatelessSession} API. This delegating
@@ -207,7 +208,8 @@ CompletionStage<T> execute(Function<Stage.Transaction, CompletionStage<T>> work)
 		 * and an error thrown by the work.
 		 */
 		CompletionStage<T> executeInTransaction(Function<Stage.Transaction, CompletionStage<T>> work) {
-			return work.apply( this )
+			return voidFuture()
+					.thenCompose( v -> work.apply( this ) )
 					// have to capture the error here and pass it along,
 					// since we can't just return a CompletionStage that
 					// rolls back the transaction from the handle() function

hibernate-reactive-core/src/test/java/org/hibernate/reactive/LazyInitializationExceptionTest.java

@@ -17,6 +17,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import io.smallrye.mutiny.Uni;
 import io.vertx.junit5.Timeout;
 import io.vertx.junit5.VertxTestContext;
 import jakarta.persistence.Entity;
@@ -32,6 +33,7 @@
 import static java.util.concurrent.TimeUnit.MINUTES;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hibernate.reactive.testing.ReactiveAssertions.assertThrown;
+import static org.hibernate.reactive.util.impl.CompletionStages.completedFuture;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -61,6 +63,31 @@ public void populateDB(VertxTestContext context) {
 		test( context, getSessionFactory().withTransaction( session -> session.persist( artemisia, liuto, sev ) ) );
 	}
 
+	@Test
+	public void testLazyInitializationExceptionWithTransactionWithMutiny(VertxTestContext context) {
+		test( context, assertThrown( LazyInitializationException.class, getMutinySessionFactory()
+					  .withSession( ms -> ms
+							  .createSelectionQuery( "from Artist", Artist.class )
+							  .getSingleResult() )
+					  .call( artist -> getMutinySessionFactory().withTransaction( s -> Uni.createFrom()
+							  // .size should throw LazyInitializationException
+							  .item( artist.getPaintings().size() ) ) )
+			  )
+		);
+	}
+
+	@Test
+	public void testLazyInitializationExceptionWithTransactionWithStage(VertxTestContext context) {
+		test( context, assertThrown( LazyInitializationException.class, getSessionFactory()
+					  .withSession( ss -> ss
+							  .createSelectionQuery( "from Artist", Artist.class )
+							  .getSingleResult() )
+					  .thenCompose( artist -> getSessionFactory()
+							  .withTransaction( s -> completedFuture( artist.getPaintings().size() ) ) )
+			  )
+		);
+	}
+
 	@Test
 	public void testLazyInitializationExceptionWithMutiny(VertxTestContext context) {
 		test( context, assertThrown( LazyInitializationException.class, openMutinySession()

hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReactiveStatelessProxyUpdateTest.java

@@ -30,6 +30,7 @@
 
 import static java.util.concurrent.TimeUnit.MINUTES;
 import static org.hibernate.reactive.testing.ReactiveAssertions.assertThrown;
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -83,26 +84,47 @@ public void testUnfetchedEntityException(VertxTestContext context) {
 	}
 
 	@Test
-	public void testLazyInitializationException(VertxTestContext context) {
+	public void testLazyInitializationExceptionWithMutiny(VertxTestContext context) {
 		Game lol = new Game( "League of Legends" );
 		GameCharacter ck = new GameCharacter( "Caitlyn Kiramman" );
 		ck.setGame( lol );
 
-		test( context, assertThrown( LazyInitializationException.class, getMutinySessionFactory()
+		test( context, getMutinySessionFactory()
 				.withTransaction( s -> s.persistAll( lol, ck ) )
 				.chain( targetId -> getMutinySessionFactory()
 						.withStatelessSession( session -> session.get( GameCharacter.class, ck.getId() ) )
 				)
-				.call( charFound -> getMutinySessionFactory()
-						.withStatelessTransaction( s -> {
+				.call( charFound -> assertThrown(
+						LazyInitializationException.class, getMutinySessionFactory().withStatelessTransaction( s -> {
 							Game game = charFound.getGame();
 							// LazyInitializationException here because we haven't fetched the entity
 							game.setGameTitle( "League of Legends V2" );
-							context.failNow( "We were expecting a LazyInitializationException" );
-							return null;
+							return Uni.createFrom().voidItem();
 						} )
+				) )
+		);
+	}
+
+	@Test
+	public void testLazyInitializationExceptionWithStage(VertxTestContext context) {
+		Game lol = new Game( "League of Legends" );
+		GameCharacter ck = new GameCharacter( "Caitlyn Kiramman" );
+		ck.setGame( lol );
+
+		test( context, getSessionFactory()
+				.withTransaction( s -> s.persist( lol, ck ) )
+				.thenCompose( targetId -> getSessionFactory()
+						.withStatelessSession( session -> session.get( GameCharacter.class, ck.getId() ) )
 				)
-		) );
+				.thenCompose( charFound -> assertThrown(
+						LazyInitializationException.class, getSessionFactory().withStatelessTransaction( s -> {
+							Game game = charFound.getGame();
+							// LazyInitializationException here because we haven't fetched the entity
+							game.setGameTitle( "League of Legends V2" );
+							return voidFuture();
+						} )
+				) )
+		);
 	}
 
 	@Test