From 0b54c1d0838ca50a1f28ffc60583e3775c790576 Mon Sep 17 00:00:00 2001
From: Andrea Boriero <andrea@hibernate.org>
Date: Wed, 12 Apr 2023 10:10:01 +0200
Subject: [PATCH] HHH-16394 Statement Batch + Version + Dirty Collection leads
 to OptimisticLockException: Batch update returned unexpected row count from
 update

---
 .../engine/jdbc/batch/internal/BatchImpl.java |  3 --
 .../mutation/UpdateCoordinatorStandard.java   | 31 ++++++++++++++++++-
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/BatchImpl.java b/hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/BatchImpl.java
index 100042d2e7..a22edbc1e9 100644
--- a/hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/BatchImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/BatchImpl.java
@@ -22,7 +22,6 @@ import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
 import org.hibernate.engine.jdbc.spi.SqlExceptionHelper;
 import org.hibernate.engine.jdbc.spi.SqlStatementLogger;
-import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.resource.jdbc.spi.JdbcObserver;
 
 import static org.hibernate.engine.jdbc.JdbcLogging.JDBC_MESSAGE_LOGGER;
@@ -108,8 +107,6 @@ public class BatchImpl implements Batch {
 			);
 		}
 
-		final SharedSessionContractImplementor session = (SharedSessionContractImplementor) jdbcCoordinator.getJdbcSessionOwner();
-
 		try {
 			getStatementGroup().forEachStatement( (tableName, statementDetails) -> {
 				if ( inclusionChecker != null && !inclusionChecker.include( statementDetails.getMutatingTableDetails() ) ) {
diff --git a/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java b/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java
index 8b5d166157..07053e9bd9 100644
--- a/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java
+++ b/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java
@@ -23,6 +23,8 @@ import org.hibernate.engine.jdbc.mutation.JdbcValueBindings;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
 import org.hibernate.engine.jdbc.mutation.ParameterUsage;
 import org.hibernate.engine.jdbc.mutation.internal.MutationQueryOptions;
+import org.hibernate.engine.jdbc.mutation.internal.NoBatchKeyAccess;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
 import org.hibernate.engine.jdbc.mutation.spi.MutationExecutorService;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -81,6 +83,7 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 	private final BatchKey batchKey;
 
 	private final MutationOperationGroup versionUpdateGroup;
+	private final BatchKey versionUpdateBatchkey;
 
 	public UpdateCoordinatorStandard(AbstractEntityPersister entityPersister, SessionFactoryImplementor factory) {
 		super( entityPersister, factory );
@@ -92,12 +95,17 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 		if ( entityPersister.hasUpdateGeneratedProperties() ) {
 			// disable batching in case of update generated properties
 			this.batchKey = null;
+			this.versionUpdateBatchkey = null;
 		}
 		else {
 			this.batchKey = new BasicBatchKey(
 					entityPersister.getEntityName() + "#UPDATE",
 					null
 			);
+			this.versionUpdateBatchkey = new BasicBatchKey(
+					entityPersister.getEntityName() + "#UPDATE_VERSION",
+					null
+			);
 		}
 	}
 
@@ -447,7 +455,7 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 
 		final EntityTableMapping mutatingTableDetails = (EntityTableMapping) versionUpdateGroup.getSingleOperation().getTableDetails();
 
-		final MutationExecutor mutationExecutor = executor( session, versionUpdateGroup, false );
+		final MutationExecutor mutationExecutor = updateVersionExecutor( session, versionUpdateGroup, false );
 
 		final EntityVersionMapping versionMapping = entityPersister().getVersionMapping();
 
@@ -973,6 +981,27 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 				.createExecutor( resolveBatchKeyAccess( dynamicUpdate, session ), group, session );
 	}
 
+	private MutationExecutor updateVersionExecutor(SharedSessionContractImplementor session, MutationOperationGroup group, boolean dynamicUpdate) {
+		return session.getSessionFactory()
+				.getServiceRegistry()
+				.getService( MutationExecutorService.class )
+				.createExecutor( resolveUpdateVersionBatchKeyAccess( dynamicUpdate, session ), group, session );
+	}
+
+	protected BatchKeyAccess resolveUpdateVersionBatchKeyAccess(boolean dynamicUpdate, SharedSessionContractImplementor session) {
+		if ( !dynamicUpdate
+				&& session.getTransactionCoordinator() != null
+				&& session.getTransactionCoordinator().isTransactionActive() ) {
+			return this::getVersionUpdateBatchkey;
+		}
+
+		return NoBatchKeyAccess.INSTANCE;
+	}
+
+	private BatchKey getVersionUpdateBatchkey(){
+		return versionUpdateBatchkey;
+	}
+
 	protected MutationOperationGroup generateDynamicUpdateGroup(
 			Object id,
 			Object rowId,