Ensure MutationExecutor triggers autoflushing before running the mutation query

Author: Sanne

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(