[hibernate#2108] StatelessSession insertAll in batch does not do batching

Author: dreab8

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

@@ -163,7 +163,7 @@ public Uni<Void> deleteAll(Object... entities) {
 
 	@Override
 	public Uni<Void> deleteAll(int batchSize, Object... entities) {
-		return uni( () -> delegate.reactiveDeleteAll( entities ) );
+		return uni( () -> delegate.reactiveDeleteAll( batchSize, entities ) );
 	}
 
 	@Override

hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/ReactiveStatelessSession.java

@@ -77,4 +77,8 @@ public interface ReactiveStatelessSession extends ReactiveQueryProducer, Reactiv
 	void close(CompletableFuture<Void> closing);
 
 	Object getIdentifier(Object entity);
+
+	void setJdbcBatchSize(Integer jdbcBatchSize);
+
+	Integer getJdbcBatchSize();
 }

hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveStatelessSessionImpl.java

@@ -150,10 +150,12 @@ private ReactiveStatelessSessionImpl(
 			PersistenceContext persistenceContext) {
 		super( factory, options );
 		this.persistenceContext = persistenceContext;
-		Integer batchSize = getConfiguredJdbcBatchSize();
-		reactiveConnection = batchSize == null || batchSize < 2
-				? connection
-				: new BatchingConnection( connection, batchSize );
+		// Don't know why ORM set the batch size to 0 in the super constructor,
+		// but the default should be the one in the factory options
+		int defaultBatchSize = factory.getSessionFactoryOptions().getJdbcBatchSize();
+		setJdbcBatchSize( defaultBatchSize );
+		// BatchingConnection will use the original connection if defaultBatchSize <= 1
+		reactiveConnection = new BatchingConnection( connection, defaultBatchSize );
 		batchingHelperSession = this;
 		influencers = new LoadQueryInfluencers( factory );
 	}
@@ -551,9 +553,12 @@ public CompletionStage<Void> reactiveInsertAll(Object... entities) {
 
 	@Override
 	public CompletionStage<Void> reactiveInsertAll(int batchSize, Object... entities) {
+		final Integer jdbcBatchSize = batchingHelperSession.getJdbcBatchSize();
+		batchingHelperSession.setJdbcBatchSize( batchSize );
 		final ReactiveConnection connection = batchingConnection( batchSize );
 		return loop( entities, batchingHelperSession::reactiveInsert )
-				.thenCompose( v -> connection.executeBatch() );
+				.thenCompose( v -> connection.executeBatch() )
+				.whenComplete( (v, throwable) -> batchingHelperSession.setJdbcBatchSize( jdbcBatchSize ) );
 	}
 
 	@Override
@@ -564,9 +569,12 @@ public CompletionStage<Void> reactiveUpdateAll(Object... entities) {
 
 	@Override
 	public CompletionStage<Void> reactiveUpdateAll(int batchSize, Object... entities) {
+		final Integer jdbcBatchSize = batchingHelperSession.getJdbcBatchSize();
+		batchingHelperSession.setJdbcBatchSize( batchSize );
 		final ReactiveConnection connection = batchingConnection( batchSize );
 		return loop( entities, batchingHelperSession::reactiveUpdate )
-				.thenCompose( v -> connection.executeBatch() );
+				.thenCompose( v -> connection.executeBatch() )
+				.whenComplete( (v, throwable) -> batchingHelperSession.setJdbcBatchSize( jdbcBatchSize ) );
 	}
 
 	@Override
@@ -577,9 +585,11 @@ public CompletionStage<Void> reactiveDeleteAll(Object... entities) {
 
 	@Override
 	public CompletionStage<Void> reactiveDeleteAll(int batchSize, Object... entities) {
+		final Integer jdbcBatchSize = batchingHelperSession.getJdbcBatchSize();
+		batchingHelperSession.setJdbcBatchSize( batchSize );
 		final ReactiveConnection connection = batchingConnection( batchSize );
-		return loop( entities, batchingHelperSession::reactiveDelete )
-				.thenCompose( v -> connection.executeBatch() );
+		return loop( entities, batchingHelperSession::reactiveDelete ).thenCompose( v -> connection.executeBatch() )
+				.whenComplete( (v, throwable) -> batchingHelperSession.setJdbcBatchSize( jdbcBatchSize ) );
 	}
 
 
@@ -591,9 +601,12 @@ public CompletionStage<Void> reactiveRefreshAll(Object... entities) {
 
 	@Override
 	public CompletionStage<Void> reactiveRefreshAll(int batchSize, Object... entities) {
+		final Integer jdbcBatchSize = batchingHelperSession.getJdbcBatchSize();
+		batchingHelperSession.setJdbcBatchSize( batchSize );
 		final ReactiveConnection connection = batchingConnection( batchSize );
 		return loop( entities, batchingHelperSession::reactiveRefresh )
-				.thenCompose( v -> connection.executeBatch() );
+				.thenCompose( v -> connection.executeBatch() )
+				.whenComplete( (v, throwable) -> batchingHelperSession.setJdbcBatchSize( jdbcBatchSize ) );
 	}
 
 	private ReactiveConnection batchingConnection(int batchSize) {