squash! HHH-19278 fix some logic in MultiIdEntityLoaders that just looks wrong

gavinking · gavinking · commit 1a613f33bd62 · 2025-03-21T11:58:34.000+01:00
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/internal/BatchFetchQueueHelper.java b/hibernate-core/src/main/java/org/hibernate/engine/internal/BatchFetchQueueHelper.java
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractMultiIdEntityLoader.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractMultiIdEntityLoader.java
@@ -8,6 +8,7 @@
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
+import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -25,7 +26,6 @@
 import java.util.ArrayList;
 import java.util.List;
 
-import static org.hibernate.engine.internal.BatchFetchQueueHelper.removeBatchLoadableEntityKey;
 import static org.hibernate.event.spi.LoadEventListener.GET;
 import static org.hibernate.internal.util.collections.CollectionHelper.arrayList;
 import static org.hibernate.internal.util.collections.CollectionHelper.isEmpty;
@@ -176,7 +176,7 @@ protected void handleResults(
 			// the element value at this position in the results List should be
 			// the EntityKey for that entity - reuse it
 			final EntityKey entityKey = (EntityKey) results.get( position );
-			removeBatchLoadableEntityKey( entityKey, session );
+			session.getPersistenceContextInternal().getBatchFetchQueue().removeBatchLoadableEntityKey( entityKey );
 			final Object entity = persistenceContext.getEntity( entityKey );
 			final Object result;
 			if ( entity == null
@@ -258,6 +258,16 @@ protected List<T> unorderedMultiLoad(
 						(position, entityKey, resolvedRef) -> result.add( (T) resolvedRef ) );
 		if ( !isEmpty( unresolvableIds ) ) {
 			loadEntitiesWithUnresolvedIds( loadOptions, lockOptions, session, unresolvableIds, result );
+			final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
+			final EntityPersister persister = getLoadable().getEntityPersister();
+			for ( Object id : unresolvableIds ) {
+				// skip any of the null padded ids
+				// (actually we could probably even break on the first null)
+				if ( id != null ) {
+					// found or not, remove the key from the batch-fetch queue
+					batchFetchQueue.removeBatchLoadableEntityKey( session.generateEntityKey( id, persister ) );
+				}
+			}
 		}
 		return result;
 	}
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java
@@ -9,6 +9,7 @@
 import java.util.Locale;
 
 import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.internal.build.AllowReflection;
@@ -17,14 +18,14 @@
 import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.JdbcMapping;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.internal.JdbcParameterImpl;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 
-import static org.hibernate.engine.internal.BatchFetchQueueHelper.removeBatchLoadableEntityKey;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.trimIdBatch;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -40,7 +41,6 @@ public class EntityBatchLoaderArrayParam<T>
 		implements SqlArrayMultiKeyLoader {
 	private final int domainBatchSize;
 
-	private final LoadQueryInfluencers loadQueryInfluencers;
 	private final BasicEntityIdentifierMapping identifierMapping;
 	private final JdbcMapping arrayJdbcMapping;
 	private final JdbcParameter jdbcParameter;
@@ -64,7 +64,6 @@ public EntityBatchLoaderArrayParam(
 			EntityMappingType entityDescriptor,
 			LoadQueryInfluencers loadQueryInfluencers) {
 		super( entityDescriptor, loadQueryInfluencers );
-		this.loadQueryInfluencers = loadQueryInfluencers;
 		this.domainBatchSize = domainBatchSize;
 
 		if ( MULTI_KEY_LOAD_LOGGER.isDebugEnabled() ) {
@@ -133,10 +132,12 @@ protected void initializeEntities(
 					getLoadable().getEntityName(), id, Arrays.toString(idsToInitialize) );
 		}
 
+		final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
+		final EntityPersister persister = getLoadable().getEntityPersister();
 		for ( Object initializedId : idsToInitialize ) {
 			if ( initializedId != null ) {
 				// found or not, remove the key from the batch-fetch queue
-				removeBatchLoadableEntityKey( initializedId, getLoadable(), session );
+				batchFetchQueue.removeBatchLoadableEntityKey( session.generateEntityKey( initializedId, persister ) );
 			}
 		}
 
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java
@@ -9,14 +9,14 @@
 
 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.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;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
@@ -41,7 +41,6 @@ public class EntityBatchLoaderInPredicate<T>
 	private final int domainBatchSize;
 	private final int sqlBatchSize;
 
-	private final LoadQueryInfluencers loadQueryInfluencers;
 	private final JdbcParametersList jdbcParameters;
 	private final SelectStatement sqlAst;
 	private final JdbcOperationQuerySelect jdbcSelectOperation;
@@ -54,7 +53,6 @@ public EntityBatchLoaderInPredicate(
 			EntityMappingType entityDescriptor,
 			LoadQueryInfluencers loadQueryInfluencers) {
 		super( entityDescriptor, loadQueryInfluencers );
-		this.loadQueryInfluencers = loadQueryInfluencers;
 		this.domainBatchSize = domainBatchSize;
 		int idColumnCount =
 				entityDescriptor.getEntityPersister().getIdentifierType()
@@ -140,6 +138,7 @@ protected void initializeEntities(
 		);
 
 		final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
+		final EntityPersister persister = getLoadable().getEntityPersister();
 
 		chunker.processChunks(
 				idsToInitialize,
@@ -164,11 +163,7 @@ protected void initializeEntities(
 				},
 				(key, relativePosition, absolutePosition) -> {
 					if ( key != null ) {
-						final EntityKey entityKey = session.generateEntityKey(
-								key,
-								getLoadable().getEntityPersister()
-						);
-						batchFetchQueue.removeBatchLoadableEntityKey( entityKey );
+						batchFetchQueue.removeBatchLoadableEntityKey( session.generateEntityKey( key, persister ) );
 					}
 				},
 				(startIndex) -> {
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiIdEntityLoaderArrayParam.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiIdEntityLoaderArrayParam.java
@@ -31,7 +31,6 @@
 import org.hibernate.sql.results.spi.ManagedResultConsumer;
 
 import static java.lang.Boolean.TRUE;
-import static org.hibernate.engine.internal.BatchFetchQueueHelper.removeBatchLoadableEntityKey;
 import static org.hibernate.engine.spi.SubselectFetch.createRegistrationHandler;
 import static org.hibernate.loader.ast.internal.LoaderHelper.loadByArrayParameter;
 import static org.hibernate.loader.ast.internal.LoaderSelectBuilder.createSelectBySingleArrayParameter;
@@ -156,15 +155,6 @@ protected void loadEntitiesWithUnresolvedIds(
 				session
 		);
 		result.addAll( databaseResults );
-
-		for ( Object id : unresolvableIds ) {
-			// skip any of the null padded ids
-			// (actually we could probably even break on the first null)
-			if ( id != null ) {
-				// found or not, remove the key from the batch-fetch queue
-				removeBatchLoadableEntityKey( id, getLoadable(), session );
-			}
-		}
 	}
 
 	@Override
