Fixing auto-flush on mutation queries

hibernate-reactive-core/src/main/java/org/hibernate/reactive/sql/exec/internal/StandardReactiveJdbcMutationExecutor.java

@@ -17,6 +17,7 @@
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.reactive.adaptor.impl.PreparedStatementAdaptor;
+import org.hibernate.reactive.engine.spi.ReactiveSharedSessionContractImplementor;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.pool.ReactiveConnection;
@@ -48,30 +49,34 @@ public CompletionStage<Integer> executeReactive(
 			Function<String, PreparedStatement> statementCreator,
 			BiConsumer<Integer, PreparedStatement> expectationCheck,
 			ExecutionContext executionContext) {
-		final SharedSessionContractImplementor session = executionContext.getSession();
-		session.autoFlushIfRequired( jdbcMutation.getAffectedTableNames() );
-
-		final LogicalConnectionImplementor logicalConnection = session
-				.getJdbcCoordinator()
-				.getLogicalConnection();
-
-		final JdbcServices jdbcServices = session.getJdbcServices();
-		final QueryOptions queryOptions = executionContext.getQueryOptions();
-		final String finalSql = finalSql( jdbcMutation, executionContext, jdbcServices, queryOptions );
-
-		Object[] parameters = PreparedStatementAdaptor
-				.bind( statement -> prepareStatement( jdbcMutation, statement, jdbcParameterBindings, executionContext ) );
-
-		session.getEventListenerManager().jdbcExecuteStatementStart();
-		return connection( executionContext )
-				.update( finalSql, parameters )
-				.thenApply( result -> {
-					// FIXME: I don't have a preparedStatement
-//					expectationCheck.accept( result, preparedStatement );
-					return result;
-				} )
-				.whenComplete( (result, t) -> session.getEventListenerManager().jdbcExecuteStatementEnd() )
-				.whenComplete( (result, t) -> executionContext.afterStatement( logicalConnection ) );
+		SharedSessionContractImplementor session = executionContext.getSession();
+		ReactiveSharedSessionContractImplementor reactiveSession = (ReactiveSharedSessionContractImplementor) session;
+
+		return reactiveSession.reactiveAutoFlushIfRequired( jdbcMutation.getAffectedTableNames() )
+				.thenCompose( v -> {
+
+			final LogicalConnectionImplementor logicalConnection = session
+					.getJdbcCoordinator()
+					.getLogicalConnection();
+
+			final JdbcServices jdbcServices = session.getJdbcServices();
+			final QueryOptions queryOptions = executionContext.getQueryOptions();
+			final String finalSql = finalSql( jdbcMutation, executionContext, jdbcServices, queryOptions );
+
+			Object[] parameters = PreparedStatementAdaptor
+					.bind( statement -> prepareStatement( jdbcMutation, statement, jdbcParameterBindings, executionContext ) );
+
+			session.getEventListenerManager().jdbcExecuteStatementStart();
+			return connection( executionContext )
+					.update( finalSql, parameters )
+					.thenApply( result -> {
+						// FIXME: I don't have a preparedStatement
+	//					expectationCheck.accept( result, preparedStatement );
+						return result;
+					} )
+					.whenComplete( (result, t) -> session.getEventListenerManager().jdbcExecuteStatementEnd() )
+					.whenComplete( (result, t) -> executionContext.afterStatement( logicalConnection ) );
+		} );
 	}
 
 	private void prepareStatement(

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

@@ -0,0 +1,114 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive;
+
+import java.util.Collection;
+import java.util.List;
+
+
+import org.junit.Test;
+
+import io.vertx.ext.unit.TestContext;
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+
+/**
+ * Verifies that scheduled changes in the reactive session
+ * are being auto-flushed before a mutation query on the same
+ * table space is executed.
+ */
+public class PersistThenDeleteTest extends BaseReactiveTest {
+
+	@Override
+	protected Collection<Class<?>> annotatedEntities() {
+		return List.of( PersistThenDeleteTest.Person.class );
+	}
+
+	@Test
+	public void testPersistThenDelete(TestContext context) {
+		test(context,
+			 getSessionFactory().withTransaction(
+							 (s, t) -> s.persist( newPerson( "foo" ), newPerson( "bar" ), newPerson( "baz" ) )
+					 )
+					 .thenCompose(
+							 v ->
+							 getSessionFactory().withTransaction(
+											 (s, t) -> s.createQuery( "from Person" ).getResultList().thenAccept( l ->
+											 context.assertEquals( 3, l.size() )
+											 )
+							 )
+					 )
+					 .thenCompose(
+							v ->
+							getSessionFactory().withTransaction(
+									(s, t) ->
+										s.persist( newPerson( "critical" ) )
+												.thenCompose( vo -> s.createQuery( "delete from Person" )
+														.executeUpdate() )
+								  )
+					 )
+					 .thenCompose(
+							 v ->
+									 getSessionFactory().withTransaction(
+											 (s, t) -> s.createQuery( "from Person" ).getResultList().thenAccept( l ->
+																														  context.assertEquals( 0, l.size() )
+											 )
+									 )
+					 )
+		);
+	}
+
+	@Test
+	public void testDeleteThenPersist(TestContext context) {
+		test(context,
+			 getSessionFactory().withTransaction(
+							 (s, t) -> s.persist( newPerson( "foo" ), newPerson( "bar" ), newPerson( "baz" ) )
+					 )
+					 .thenCompose(
+							 v ->
+									 getSessionFactory().withTransaction(
+											 (s, t) -> s.createQuery( "from Person" ).getResultList().thenAccept( l ->
+																														  context.assertEquals( 3, l.size() )
+											 )
+									 )
+					 )
+					 .thenCompose(
+							 v ->
+									 getSessionFactory().withTransaction(
+											 (s, t) ->
+													 s.createQuery( "delete from Person" ).executeUpdate()
+															 .thenCompose( vo -> s.persist( newPerson( "critical" ) ) )
+									 )
+					 )
+					 .thenCompose(
+							 v ->
+									 getSessionFactory().withTransaction(
+											 (s, t) -> s.createQuery( "from Person" ).getResultList().thenAccept( l ->
+																														  context.assertEquals( 1, l.size() )
+											 )
+									 )
+					 )
+		);
+	}
+
+	private static Person newPerson(String name) {
+		final Person person = new Person();
+		person.name = name;
+		return person;
+	}
+
+	@Entity(name="Person")
+	public static class Person {
+
+		@Id
+		@GeneratedValue
+		Integer id;
+
+		String name;
+	}
+
+}