HHH-16390 - Execution of non-batched statements do not force execution of current batch

HHH-16319 - test

hibernate-core/src/main/java/org/hibernate/engine/jdbc/internal/JdbcCoordinatorImpl.java

@@ -18,6 +18,7 @@ import java.util.function.Supplier;
 import org.hibernate.ConnectionReleaseMode;
 import org.hibernate.HibernateException;
 import org.hibernate.TransactionException;
+import org.hibernate.engine.jdbc.batch.JdbcBatchLogging;
 import org.hibernate.engine.jdbc.batch.spi.Batch;
 import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.group.PreparedStatementGroup;
@@ -192,6 +193,18 @@ public class JdbcCoordinatorImpl implements JdbcCoordinator {
 		}
 	}
 
+	@Override
+	public void conditionallyExecuteBatch(BatchKey key) {
+		if ( currentBatch == null ) {
+			return;
+		}
+
+		if ( !currentBatch.getKey().equals( key ) ) {
+			JdbcBatchLogging.BATCH_LOGGER.debugf( "Conditionally executing batch - %s", currentBatch.getKey() );
+			currentBatch.execute();
+		}
+	}
+
 	@Override
 	public void abortBatch() {
 		if ( currentBatch != null ) {

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/AbstractMutationExecutor.java

@@ -8,6 +8,7 @@ package org.hibernate.engine.jdbc.mutation.internal;
 
 import java.sql.SQLException;
 
+import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.JdbcValueBindings;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
 import org.hibernate.engine.jdbc.mutation.OperationResultChecker;
@@ -26,6 +27,15 @@ import static org.hibernate.sql.model.ModelMutationLogging.MODEL_MUTATION_LOGGER
  * @author Steve Ebersole
  */
 public abstract class AbstractMutationExecutor implements MutationExecutor {
+	/**
+	 * Executors with non-batched operations should call this to clean up any "previous" batch
+	 * before starting their work
+	 */
+	protected void prepareForNonBatchedWork(BatchKey batchKey, SharedSessionContractImplementor session) {
+		// if there is a current batch, make sure to execute it first
+		session.getJdbcCoordinator().conditionallyExecuteBatch( batchKey );
+	}
+
 	/**
 	 * Templated implementation of execution as <ol>
 	 *     <li>{@link #performNonBatchedOperations}</li>

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/MutationExecutorSingleNonBatched.java

@@ -22,8 +22,8 @@ public class MutationExecutorSingleNonBatched extends AbstractSingleMutationExec
 			PreparableMutationOperation mutationOperation,
 			SharedSessionContractImplementor session) {
 		super( mutationOperation, session );
-
 		this.statementGroup = new PreparedStatementGroupSingleTable( mutationOperation, session );
+		prepareForNonBatchedWork( null, session );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/MutationExecutorSingleSelfExecuting.java

@@ -33,6 +33,8 @@ public class MutationExecutorSingleSelfExecuting extends AbstractMutationExecuto
 				this::findJdbcValueDescriptor,
 				session
 		);
+
+		prepareForNonBatchedWork( null, session );
 	}
 
 	private JdbcValueDescriptor findJdbcValueDescriptor(String tableName, String columnName, ParameterUsage usage) {

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/MutationExecutorStandard.java

@@ -21,7 +21,9 @@ import org.hibernate.engine.jdbc.mutation.ParameterUsage;
 import org.hibernate.engine.jdbc.mutation.TableInclusionChecker;
 import org.hibernate.engine.jdbc.mutation.group.PreparedStatementDetails;
 import org.hibernate.engine.jdbc.mutation.group.PreparedStatementGroup;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.sql.model.MutationOperation;
 import org.hibernate.sql.model.MutationOperationGroup;
 import org.hibernate.sql.model.PreparableMutationOperation;
@@ -30,6 +32,8 @@ import org.hibernate.sql.model.TableMapping;
 import org.hibernate.sql.model.ValuesAnalysis;
 import org.hibernate.sql.model.jdbc.JdbcValueDescriptor;
 
+import static org.hibernate.internal.util.collections.CollectionHelper.isNotEmpty;
+
 /**
  * Standard MutationExecutor implementation
  *
@@ -60,12 +64,12 @@ public class MutationExecutorStandard extends AbstractMutationExecutor {
 
 	public MutationExecutorStandard(
 			MutationOperationGroup mutationOperationGroup,
-			Supplier<BatchKey> batchKeySupplier,
+			BatchKeyAccess batchKeySupplier,
 			int batchSize,
 			SharedSessionContractImplementor session) {
 		this.mutationOperationGroup = mutationOperationGroup;
 
-		final BatchKey batchKey = batchKeySupplier.get();
+		final BatchKey batchKey = batchKeySupplier.getBatchKey();
 
 		// split the table operations into batchable and non-batchable -
 		// 		1. batchable statements are handle via Batch
@@ -155,6 +159,10 @@ public class MutationExecutorStandard extends AbstractMutationExecutor {
 				this::findJdbcValueDescriptor,
 				session
 		);
+
+		if ( isNotEmpty( nonBatchedJdbcMutations ) || isNotEmpty( selfExecutingMutations ) ) {
+			prepareForNonBatchedWork( batchKey, session );
+		}
 	}
 
 	protected PreparedStatementGroup getNonBatchedStatementGroup() {

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/NoBatchKeyAccess.java

@@ -0,0 +1,28 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html.
+ */
+package org.hibernate.engine.jdbc.mutation.internal;
+
+import org.hibernate.engine.jdbc.batch.spi.BatchKey;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
+
+/**
+ * A form of BatchKeyAccess for cases where batching is not wanted, which is
+ * signified by a BatchKey of {@code null}
+ *
+ * @author Steve Ebersole
+ */
+public class NoBatchKeyAccess implements BatchKeyAccess {
+	/**
+	 * Singleton access
+	 */
+	public static final NoBatchKeyAccess INSTANCE = new NoBatchKeyAccess();
+
+	@Override
+	public BatchKey getBatchKey() {
+		return null;
+	}
+}

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/internal/StandardMutationExecutorService.java

@@ -12,6 +12,7 @@ import java.util.function.Supplier;
 import org.hibernate.cfg.Environment;
 import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
 import org.hibernate.engine.jdbc.mutation.spi.MutationExecutorService;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.internal.util.config.ConfigurationHelper;
@@ -44,7 +45,7 @@ public class StandardMutationExecutorService implements MutationExecutorService
 
 	@Override
 	public MutationExecutor createExecutor(
-			Supplier<BatchKey> batchKeySupplier,
+			BatchKeyAccess batchKeySupplier,
 			MutationOperationGroup operationGroup,
 			SharedSessionContractImplementor session) {
 		// decide whether to use batching - any number > one means to batch
@@ -78,7 +79,7 @@ public class StandardMutationExecutorService implements MutationExecutorService
 			}
 
 			final PreparableMutationOperation jdbcOperation = (PreparableMutationOperation) singleOperation;
-			final BatchKey batchKey = batchKeySupplier.get();
+			final BatchKey batchKey = batchKeySupplier.getBatchKey();
 			if ( jdbcOperation.canBeBatched( batchKey, batchSizeToUse ) ) {
 				return new MutationExecutorSingleBatched( jdbcOperation, batchKey, batchSizeToUse, session );
 			}

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/spi/BatchKeyAccess.java

@@ -0,0 +1,21 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html.
+ */
+package org.hibernate.engine.jdbc.mutation.spi;
+
+import org.hibernate.engine.jdbc.batch.spi.BatchKey;
+
+/**
+ * Provides access to a BatchKey as part of creating an {@linkplain MutationExecutorService#createExecutor executor}.
+ *
+ * @author Steve Ebersole
+ */
+public interface BatchKeyAccess {
+	/**
+	 * The BatchKey to use
+	 */
+	BatchKey getBatchKey();
+}

hibernate-core/src/main/java/org/hibernate/engine/jdbc/mutation/spi/MutationExecutorService.java

@@ -6,9 +6,6 @@
  */
 package org.hibernate.engine.jdbc.mutation.spi;
 
-import java.util.function.Supplier;
-
-import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.service.Service;
@@ -21,8 +18,11 @@ import org.hibernate.sql.model.MutationOperationGroup;
  */
 public interface MutationExecutorService extends Service {
 
+	/**
+	 * Create an executor for the given {@code operationGroup}, potentially using batching
+	 */
 	MutationExecutor createExecutor(
-			Supplier<BatchKey> batchKeySupplier,
+			BatchKeyAccess batchKeySupplier,
 			MutationOperationGroup operationGroup,
 			SharedSessionContractImplementor session);
 }

hibernate-core/src/main/java/org/hibernate/engine/jdbc/spi/JdbcCoordinator.java

@@ -60,6 +60,12 @@ public interface JdbcCoordinator extends Serializable, TransactionCoordinatorOwn
 	 */
 	void executeBatch();
 
+	/**
+	 * Conditionally execute the currently managed batch (if any), if the
+	 * keys do not match
+	 */
+	void conditionallyExecuteBatch(BatchKey key);
+
 	/**
 	 * Abort the currently managed batch (if any)
 	 */

hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/AbstractMutationCoordinator.java

@@ -10,8 +10,11 @@ import java.util.List;
 
 import org.hibernate.Internal;
 import org.hibernate.dialect.Dialect;
+import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.JdbcValueBindings;
 import org.hibernate.engine.jdbc.mutation.ParameterUsage;
+import org.hibernate.engine.jdbc.mutation.internal.NoBatchKeyAccess;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.generator.OnExecutionGenerator;
@@ -60,6 +63,18 @@ public abstract class AbstractMutationCoordinator {
 		return factory().getJdbcServices().getDialect();
 	}
 
+	protected BatchKeyAccess resolveBatchKeyAccess(boolean dynamicUpdate, SharedSessionContractImplementor session) {
+		if ( !dynamicUpdate
+				&& session.getTransactionCoordinator() != null
+				&& session.getTransactionCoordinator().isTransactionActive() ) {
+			return this::getBatchKey;
+		}
+
+		return NoBatchKeyAccess.INSTANCE;
+	}
+
+	protected abstract BatchKey getBatchKey();
+
 	protected MutationOperationGroup createOperationGroup(ValuesAnalysis valuesAnalysis, MutationGroup mutationGroup) {
 		final int numberOfTableMutations = mutationGroup.getNumberOfTableMutations();
 		switch ( numberOfTableMutations ) {

hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/DeleteCoordinator.java

@@ -133,9 +133,7 @@ public class DeleteCoordinator extends AbstractMutationCoordinator {
 		return session.getFactory()
 				.getServiceRegistry()
 				.getService( MutationExecutorService.class )
-				.createExecutor( ( session.getTransactionCoordinator() != null &&
+				.createExecutor( resolveBatchKeyAccess( false, session ), group, session );
-										session.getTransactionCoordinator().isTransactionActive() ? () -> batchKey : () -> null ),
-								group, session );
 	}
 
 	protected void applyLocking(

hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/InsertCoordinator.java

@@ -12,6 +12,7 @@ import java.util.List;
 import org.hibernate.Internal;
 import org.hibernate.dialect.Dialect;
 import org.hibernate.engine.jdbc.batch.internal.BasicBatchKey;
+import org.hibernate.engine.jdbc.batch.spi.BatchKey;
 import org.hibernate.engine.jdbc.mutation.JdbcValueBindings;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
 import org.hibernate.engine.jdbc.mutation.ParameterUsage;
@@ -79,7 +80,8 @@ public class InsertCoordinator extends AbstractMutationCoordinator {
 		return staticInsertGroup;
 	}
 
-	public BasicBatchKey getInsertBatchKey() {
+	@Override
+	protected BatchKey getBatchKey() {
 		return batchKey;
 	}
 
@@ -305,9 +307,7 @@ public class InsertCoordinator extends AbstractMutationCoordinator {
 		return session.getFactory()
 				.getServiceRegistry()
 				.getService( MutationExecutorService.class )
-				.createExecutor( ( !dynamicUpdate && session.getTransactionCoordinator() != null &&
+				.createExecutor( resolveBatchKeyAccess( dynamicUpdate, session ), group, session );
-									session.getTransactionCoordinator().isTransactionActive() ? () -> batchKey : () -> null ),
-								group, session );
 	}
 
 	protected static TableInclusionChecker getTableInclusionChecker(InsertValuesAnalysis insertValuesAnalysis) {
@@ -410,4 +410,12 @@ public class InsertCoordinator extends AbstractMutationCoordinator {
 			&& generator.generatedOnExecution()
 			&& ( (OnExecutionGenerator) generator ).referenceColumnsInSql(dialect);
 	}
+
+	/**
+	 * @deprecated Use {@link #getBatchKey()}
+	 */
+	@Deprecated
+	public BasicBatchKey getInsertBatchKey() {
+		return batchKey;
+	}
 }

hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java

@@ -970,9 +970,7 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 		return session.getSessionFactory()
 				.getServiceRegistry()
 				.getService( MutationExecutorService.class )
-				.createExecutor( ( !dynamicUpdate && session.getTransactionCoordinator() != null &&
+				.createExecutor( resolveBatchKeyAccess( dynamicUpdate, session ), group, session );
-										session.getTransactionCoordinator().isTransactionActive() ? () -> batchKey : () -> null ),
-								group, session );
 	}
 
 	protected MutationOperationGroup generateDynamicUpdateGroup(

hibernate-core/src/test/java/org/hibernate/orm/test/batch/BatchGeneratedAssociationTest.java

@@ -0,0 +1,109 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.orm.test.batch;
+
+import java.io.Serializable;
+
+import org.hibernate.annotations.Generated;
+import org.hibernate.cfg.AvailableSettings;
+
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Embeddable;
+import jakarta.persistence.EmbeddedId;
+import jakarta.persistence.Entity;
+import jakarta.persistence.FetchType;
+import jakarta.persistence.Id;
+import jakarta.persistence.JoinColumn;
+import jakarta.persistence.JoinColumns;
+import jakarta.persistence.OneToOne;
+import jakarta.persistence.Table;
+
+@SessionFactory
+@DomainModel( annotatedClasses = {
+		BatchGeneratedAssociationTest.Interpretation.class,
+		BatchGeneratedAssociationTest.InterpretationData.class,
+		BatchGeneratedAssociationTest.InterpretationVersion.class
+} )
+@ServiceRegistry( settings = @Setting( name = AvailableSettings.STATEMENT_BATCH_SIZE, value = "10" ) )
+public class BatchGeneratedAssociationTest {
+	@Test
+	public void test(SessionFactoryScope scope) {
+		scope.inTransaction( session -> {
+			final Long interpretationVersion = 111L;
+
+			final Interpretation interpretation = new Interpretation();
+			interpretation.uuid = 1L;
+
+			final InterpretationData interpretationData = new InterpretationData();
+			interpretationData.interpretationVersion = new InterpretationVersion(
+					interpretationVersion,
+					interpretation.uuid
+			);
+			interpretationData.name = "TEST_NAME";
+			session.persist( interpretationData );
+
+			interpretation.interpretationData = interpretationData;
+			interpretation.interpretationVersion = interpretationVersion;
+			session.persist( interpretation );
+		} );
+	}
+
+	@Entity( name = "Interpretation" )
+	@Table( name = "interpretations" )
+	public static class Interpretation {
+		@Id
+		public Long uuid;
+
+		@Column( name = "interpretation_version" )
+		public Long interpretationVersion;
+
+		@Column( name = "id" )
+		@Generated
+		public Long id;
+
+		@OneToOne( fetch = FetchType.LAZY )
+		@JoinColumns( {
+				@JoinColumn( name = "uuid", referencedColumnName = "interpretation_uuid", insertable = false, updatable = false ),
+				@JoinColumn( name = "interpretation_version", referencedColumnName = "interpretation_version", insertable = false, updatable = false )
+		} )
+		public InterpretationData interpretationData;
+	}
+
+	@Embeddable
+	public static class InterpretationVersion implements Serializable {
+		@Column( name = "interpretation_version", nullable = false, updatable = false )
+		public Long version;
+
+		@Column( name = "interpretation_uuid", nullable = false, updatable = false )
+		public Long uuid;
+
+		public InterpretationVersion() {
+		}
+
+		public InterpretationVersion(Long version, Long uuid) {
+			this.version = version;
+			this.uuid = uuid;
+		}
+	}
+
+	@Entity( name = "InterpretationData" )
+	@Table( name = "interpretation_data" )
+	public static class InterpretationData {
+		@EmbeddedId
+		public InterpretationVersion interpretationVersion;
+
+		@Column( updatable = false )
+		public String name;
+	}
+}