HHH-16820 When batching enabled the LockModeType is ignored

hibernate · Jun 27, 2023 · 50fab5b · 50fab5b
1 parent d6978f5
commit 50fab5b
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 41 deletions.
diff --git a/...rnate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractEntityBatchLoader.java b/...rnate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractEntityBatchLoader.java
@@ -0,0 +1,42 @@
+/*
+ * 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.loader.ast.internal;
+
+import org.hibernate.Hibernate;
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.loader.ast.spi.EntityBatchLoader;
+import org.hibernate.metamodel.mapping.EntityMappingType;
+
+public abstract class AbstractEntityBatchLoader<T>
+		extends SingleIdEntityLoaderSupport<T>
+		implements EntityBatchLoader<T>{
+
+	protected final SingleIdEntityLoaderStandardImpl<T> singleIdLoader;
+
+	public AbstractEntityBatchLoader(EntityMappingType entityDescriptor, SessionFactoryImplementor sessionFactory) {
+		super( entityDescriptor, sessionFactory );
+		singleIdLoader = new SingleIdEntityLoaderStandardImpl<>( entityDescriptor, sessionFactory );
+	}
+
+	@Override
+	public T load(
+			Object pkValue,
+			Object entityInstance,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session) {
+		final T entity = load( pkValue, entityInstance, lockOptions, null, session );
+		if ( Hibernate.isInitialized( entity ) ) {
+			return entity;
+		}
+		else {
+			return null;
+		}
+	}
+
+}
diff --git a/...ate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java b/...ate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java
@@ -9,14 +9,13 @@
 import java.lang.reflect.Array;
 import java.util.Locale;
 
-import org.hibernate.Hibernate;
+import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
-import org.hibernate.loader.ast.spi.EntityBatchLoader;
 import org.hibernate.loader.ast.spi.SqlArrayMultiKeyLoader;
 import org.hibernate.metamodel.mapping.BasicEntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
@@ -29,6 +28,8 @@
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.trimIdBatch;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -40,8 +41,8 @@
  * @author Steve Ebersole
  */
 public class EntityBatchLoaderArrayParam<T>
-		extends SingleIdEntityLoaderSupport<T>
-		implements EntityBatchLoader<T>, SqlArrayMultiKeyLoader, Preparable {
+		extends AbstractEntityBatchLoader<T>
+		implements SqlArrayMultiKeyLoader, Preparable {
 	private final int domainBatchSize;
 
 	private BasicEntityIdentifierMapping identifierMapping;
@@ -94,29 +95,18 @@ public final T load(
 			MULTI_KEY_LOAD_LOGGER.debugf( "Batch fetching entity `%s#%s`", getLoadable().getEntityName(), pkValue );
 		}
 
-		Object[] ids = resolveIdsToInitialize( pkValue, session );
+		final Object[] ids = resolveIdsToInitialize( pkValue, session );
+		if ( hasSingleId( ids ) || lockOptions.getLockMode() != LockMode.NONE ) {
+			return singleIdLoader.load( pkValue, entityInstance, lockOptions, readOnly, session );
+		}
+
 		initializeEntities( ids, pkValue, entityInstance, lockOptions, readOnly, session );
 
 		final EntityKey entityKey = session.generateEntityKey( pkValue, getLoadable().getEntityPersister() );
 		//noinspection unchecked
 		return (T) session.getPersistenceContext().getEntity( entityKey );
 	}
 
-	@Override
-	public T load(
-			Object pkValue,
-			Object entityInstance,
-			LockOptions lockOptions,
-			SharedSessionContractImplementor session) {
-		final T entity = load( pkValue, entityInstance, lockOptions, null, session );
-		if ( Hibernate.isInitialized( entity ) ) {
-			return entity;
-		}
-		else {
-			return null;
-		}
-	}
-
 	@SuppressWarnings( "unchecked" )
 	protected <X> X[] resolveIdsToInitialize(Object pkValue, SharedSessionContractImplementor session) {
 		final X[] idsToLoad = (X[]) Array.newInstance( identifierMapping.getJdbcMapping().getJdbcJavaType().getJavaTypeClass(), domainBatchSize );
@@ -128,7 +118,7 @@ protected <X> X[] resolveIdsToInitialize(Object pkValue, SharedSessionContractIm
 				pkValue,
 				getLoadable()
 		);
-		return idsToLoad;
+		return (X[]) trimIdBatch( domainBatchSize, idsToLoad );
 	}
 
 	private void initializeEntities(

diff --git a/...te-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java b/...te-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java
@@ -10,15 +10,14 @@
 import java.util.List;
 import java.util.Locale;
 
-import org.hibernate.Hibernate;
+import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.SubselectFetch;
-import org.hibernate.loader.ast.spi.EntityBatchLoader;
 import org.hibernate.loader.ast.spi.SqlInPredicateMultiKeyLoader;
 import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityMappingType;
@@ -32,6 +31,7 @@
 import org.hibernate.sql.results.spi.ListResultsConsumer;
 
 import static org.hibernate.internal.util.collections.CollectionHelper.arrayList;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -46,8 +46,8 @@
  * @author Steve Ebersole
  */
 public class EntityBatchLoaderInPredicate<T>
-		extends SingleIdEntityLoaderSupport<T>
-		implements EntityBatchLoader<T>, SqlInPredicateMultiKeyLoader, Preparable {
+		extends AbstractEntityBatchLoader<T>
+		implements SqlInPredicateMultiKeyLoader, Preparable {
 	private final int domainBatchSize;
 	private final int sqlBatchSize;
 
@@ -104,6 +104,9 @@ public final T load(
 		}
 
 		final Object[] idsToInitialize = resolveIdsToLoad( pkValue, session );
+		if ( hasSingleId( idsToInitialize ) || lockOptions.getLockMode() != LockMode.NONE ) {
+			return singleIdLoader.load( pkValue, entityInstance, lockOptions, readOnly, session );
+		}
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
 			MULTI_KEY_LOAD_LOGGER.debugf( "Ids to batch-fetch initialize (`%s#%s`) %s", getLoadable().getEntityName(), pkValue, Arrays.toString(idsToInitialize) );
 		}
@@ -115,21 +118,6 @@ public final T load(
 		return (T) session.getPersistenceContext().getEntity( entityKey );
 	}
 
-	@Override
-	public T load(
-			Object pkValue,
-			Object entityInstance,
-			LockOptions lockOptions,
-			SharedSessionContractImplementor session) {
-		final T entity = load( pkValue, entityInstance, lockOptions, null, session );
-		if ( Hibernate.isInitialized( entity ) ) {
-			return entity;
-		}
-		else {
-			return null;
-		}
-	}
-
 	protected Object[] resolveIdsToLoad(Object pkValue, SharedSessionContractImplementor session) {
 		return session.getPersistenceContextInternal().getBatchFetchQueue().getBatchLoadableEntityIds(
 				getLoadable(),

diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiKeyLoadHelper.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiKeyLoadHelper.java
@@ -6,6 +6,8 @@
  */
 package org.hibernate.loader.ast.internal;
 
+import java.util.Arrays;
+
 import org.hibernate.dialect.Dialect;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.metamodel.mapping.JdbcMapping;
@@ -56,4 +58,21 @@ public static JdbcMapping resolveArrayJdbcMapping(
 				typeConfiguration.getCurrentBaseSqlTypeIndicators()
 		);
 	}
+
+	static boolean hasSingleId(Object[] ids) {
+		for ( int i=1; i<ids.length; i++ ) {
+			if ( ids[i] != null ) {
+				return false;
+			}
+		}
+		return true;
+	}
+	static Object[] trimIdBatch(int length, Object[] keysToInitialize) {
+		int newLength = length;
+		while ( newLength>1 && keysToInitialize[newLength-1] == null ) {
+			newLength--;
+		}
+		return newLength < length ? Arrays.copyOf( keysToInitialize, newLength) : keysToInitialize;
+	}
+
 }
diff --git a/...nate-core/src/test/java/org/hibernate/orm/test/mapping/fetch/depth/DepthOneBatchTest.java b/...nate-core/src/test/java/org/hibernate/orm/test/mapping/fetch/depth/DepthOneBatchTest.java
@@ -92,7 +92,7 @@ public void tesGetAgency(SessionFactoryScope scope) {
 							"select a1_0.agency_id,a1_0.agency_txt from agency_table a1_0 where a1_0.agency_id=?"
 					);
 					assertThat( executedQueries.get( 1 ).toLowerCase() ).isEqualTo(
-							"select a1_0.agency_id,a1_0.agency_detail from agency_detail_table a1_0 where array_contains(?,a1_0.agency_id)"
+							"select a1_0.agency_id,a1_0.agency_detail from agency_detail_table a1_0 where a1_0.agency_id=?"
 					);
 
 					assertThat( executedQueries.get( 2 ).toLowerCase() ).isEqualTo(