HHH-19248 Return of deleted entities doesn't work with naturalid multiloading

hibernate-core/src/main/java/org/hibernate/action/internal/EntityDeleteAction.java

@@ -60,16 +60,6 @@ public EntityDeleteAction(
 		this.version = version;
 		this.isCascadeDeleteEnabled = isCascadeDeleteEnabled;
 		this.state = state;
-
-		final NaturalIdMapping naturalIdMapping = persister.getNaturalIdMapping();
-		if ( naturalIdMapping != null ) {
-			naturalIdValues = session.getPersistenceContextInternal().getNaturalIdResolutions()
-					.removeLocalResolution(
-							getId(),
-							naturalIdMapping.extractNaturalIdFromEntityState( state ),
-							getPersister()
-					);
-		}
 	}
 
 	/**
@@ -125,6 +115,16 @@ public void execute() throws HibernateException {
 
 		final boolean veto = isInstanceLoaded() && preDelete();
 
+		final NaturalIdMapping naturalIdMapping = persister.getNaturalIdMapping();
+		if ( naturalIdMapping != null ) {
+			naturalIdValues = session.getPersistenceContextInternal().getNaturalIdResolutions()
+					.removeLocalResolution(
+							getId(),
+							naturalIdMapping.extractNaturalIdFromEntityState( state ),
+							getPersister()
+					);
+		}
+
 		final Object ck = lockCacheItem();
 
 		if ( !isCascadeDeleteEnabled && !veto ) {

hibernate-core/src/main/java/org/hibernate/generator/OnExecutionGenerator.java

@@ -14,7 +14,7 @@
 import org.hibernate.persister.entity.EntityPersister;
 
 import static org.hibernate.generator.EventType.INSERT;
-import static org.hibernate.generator.internal.NaturalIdHelper.getNaturalIdPropertyNames;
+import static org.hibernate.internal.NaturalIdHelper.getNaturalIdPropertyNames;
 import static org.hibernate.generator.values.internal.GeneratedValuesHelper.noCustomSql;
 
 /**

hibernate-core/src/main/java/org/hibernate/generator/values/internal/GeneratedValuesHelper.java

@@ -52,7 +52,7 @@
 import org.hibernate.sql.results.spi.ListResultsConsumer;
 import org.hibernate.sql.results.spi.RowReader;
 
-import static org.hibernate.generator.internal.NaturalIdHelper.getNaturalIdPropertyNames;
+import static org.hibernate.internal.NaturalIdHelper.getNaturalIdPropertyNames;
 import static org.hibernate.pretty.MessageHelper.infoString;
 
 /**

hibernate-core/src/main/java/org/hibernate/id/IdentityGenerator.java

@@ -18,7 +18,7 @@
 import org.hibernate.persister.entity.EntityPersister;
 
 import static org.hibernate.generator.EventType.INSERT;
-import static org.hibernate.generator.internal.NaturalIdHelper.getNaturalIdPropertyNames;
+import static org.hibernate.internal.NaturalIdHelper.getNaturalIdPropertyNames;
 import static org.hibernate.generator.values.internal.GeneratedValuesHelper.noCustomSql;
 
 /**

hibernate-core/src/main/java/org/hibernate/id/SelectGenerator.java

@@ -11,7 +11,7 @@
 import org.hibernate.generator.OnExecutionGenerator;
 import org.hibernate.persister.entity.EntityPersister;
 
-import static org.hibernate.generator.internal.NaturalIdHelper.getNaturalIdPropertyNames;
+import static org.hibernate.internal.NaturalIdHelper.getNaturalIdPropertyNames;
 
 /**
  * A generator that {@code select}s the just-{@code insert}ed row to determine the

hibernate-core/src/main/java/org/hibernate/internal/NaturalIdHelper.java

@@ -0,0 +1,100 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.internal;
+
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistenceContext;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.engine.spi.Status;
+import org.hibernate.id.IdentifierGenerationException;
+import org.hibernate.loader.LoaderLogging;
+import org.hibernate.metamodel.mapping.EntityMappingType;
+import org.hibernate.metamodel.mapping.NaturalIdMapping;
+import org.hibernate.persister.entity.EntityPersister;
+
+import java.util.Collection;
+
+/**
+ * @author Gavin King
+ */
+public class NaturalIdHelper {
+	public static String[] getNaturalIdPropertyNames(EntityPersister persister) {
+		final int[] naturalIdPropertyIndices = persister.getNaturalIdentifierProperties();
+		if ( naturalIdPropertyIndices == null ) {
+			throw new IdentifierGenerationException( "Entity '" + persister.getEntityName()
+					+ "' has no '@NaturalId' property" );
+		}
+		if ( persister.getEntityMetamodel().isNaturalIdentifierInsertGenerated() ) {
+			throw new IdentifierGenerationException( "Entity '" + persister.getEntityName()
+					+ "' has a '@NaturalId' property which is also defined as insert-generated" );
+		}
+		final String[] allPropertyNames = persister.getPropertyNames();
+		final String[] propertyNames = new String[naturalIdPropertyIndices.length];
+		for ( int i = 0; i < naturalIdPropertyIndices.length; i++ ) {
+			propertyNames[i] = allPropertyNames[naturalIdPropertyIndices[i]];
+		}
+		return propertyNames;
+	}
+
+	public static void performAnyNeededCrossReferenceSynchronizations(
+			boolean synchronizationEnabled,
+			EntityMappingType entityMappingType,
+			SharedSessionContractImplementor session) {
+
+		if ( !synchronizationEnabled ) {
+			// synchronization (this process) was disabled
+			return;
+		}
+
+		final NaturalIdMapping naturalIdMapping = entityMappingType.getNaturalIdMapping();
+
+		if ( !naturalIdMapping.isMutable() ) {
+			// only mutable natural-ids need this processing
+			return;
+		}
+
+		if ( ! session.isTransactionInProgress() ) {
+			// not in a transaction so skip synchronization
+			return;
+		}
+
+		final EntityPersister entityPersister = entityMappingType.getEntityPersister();
+
+		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
+		final Collection<?> cachedPkResolutions =
+				persistenceContext.getNaturalIdResolutions()
+						.getCachedPkResolutions( entityPersister );
+		final boolean loggerDebugEnabled = LoaderLogging.LOADER_LOGGER.isDebugEnabled();
+		for ( Object pk : cachedPkResolutions ) {
+			final EntityKey entityKey = session.generateEntityKey( pk, entityPersister );
+			final Object entity = persistenceContext.getEntity( entityKey );
+			final EntityEntry entry = persistenceContext.getEntry( entity );
+
+			if ( entry == null ) {
+				if ( loggerDebugEnabled ) {
+					LoaderLogging.LOADER_LOGGER.debugf(
+							"Cached natural-id/pk resolution linked to null EntityEntry in persistence context : %s#%s",
+							entityMappingType.getEntityName(),
+							pk
+					);
+				}
+				continue;
+			}
+
+			if ( !entry.requiresDirtyCheck( entity ) ) {
+				continue;
+			}
+
+			// MANAGED is the only status we care about here...
+			if ( entry.getStatus() != Status.MANAGED ) {
+				continue;
+			}
+
+			persistenceContext.getNaturalIdResolutions().handleSynchronization( pk, entity, entityPersister );
+		}
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/internal/NaturalIdMultiLoadAccessStandard.java

@@ -5,7 +5,6 @@
 package org.hibernate.internal;
 
 import java.util.List;
-import java.util.Set;
 
 import org.hibernate.CacheMode;
 import org.hibernate.LockOptions;
@@ -16,10 +15,11 @@
 import org.hibernate.graph.GraphSemantic;
 import org.hibernate.graph.RootGraph;
 import org.hibernate.graph.spi.RootGraphImplementor;
-import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.loader.ast.spi.MultiNaturalIdLoadOptions;
 import org.hibernate.persister.entity.EntityPersister;
 
+import static org.hibernate.internal.NaturalIdHelper.performAnyNeededCrossReferenceSynchronizations;
+
 /**
  * @author Steve Ebersole
  */
@@ -82,6 +82,8 @@ public NaturalIdMultiLoadAccess<T> enableOrderedReturn(boolean enabled) {
 	@Override
 	@SuppressWarnings( "unchecked" )
 	public List<T> multiLoad(Object... ids) {
+		performAnyNeededCrossReferenceSynchronizations( true, entityDescriptor, session );
+
 		final CacheMode sessionCacheMode = session.getCacheMode();
 		boolean cacheModeChanged = false;
 
@@ -110,7 +112,6 @@ public List<T> multiLoad(Object... ids) {
 			}
 
 			try {
-				session.autoFlushIfRequired( (Set) CollectionHelper.setOf( entityDescriptor.getQuerySpaces() ) );
 				return (List<T>) entityDescriptor.getMultiNaturalIdLoader().multiLoad( ids, this, session );
 			}
 			finally {

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractMultiNaturalIdLoader.java

@@ -14,10 +14,13 @@
 import org.hibernate.loader.ast.spi.MultiNaturalIdLoader;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
+import static org.hibernate.internal.util.collections.CollectionHelper.arrayList;
+import static org.hibernate.internal.util.collections.CollectionHelper.isEmpty;
+import static org.hibernate.loader.ast.internal.LoaderHelper.upgradeLock;
+
 /**
  * @author Jan Schatteman
  */
@@ -57,7 +60,18 @@ private <K> List<E> performUnorderedMultiLoad(K[] naturalIds, MultiNaturalIdLoad
 		);
 	}
 
-	protected abstract <K> List<E> unorderedMultiLoad(K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions);
+	protected <K> List<E> unorderedMultiLoad(K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions) {
+		final List<E> results = arrayList( naturalIds.length );
+		final Object[] unresolvedIds =
+				checkPersistenceContextForCachedResults( naturalIds, session, lockOptions, results );
+		if ( !isEmpty( unresolvedIds ) ) {
+			results.addAll( loadEntitiesWithUnresolvedIds(unresolvedIds, session, lockOptions) );
+		}
+
+		return results;
+	}
+
+	protected abstract List<E> loadEntitiesWithUnresolvedIds(Object[] unresolvedIds, SharedSessionContractImplementor session, LockOptions lockOptions);
 
 	private <K> List<E> performOrderedMultiLoad(K[] naturalIds, MultiNaturalIdLoadOptions options, SharedSessionContractImplementor session) {
 		if ( MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER.isTraceEnabled() ) {
@@ -75,19 +89,18 @@ protected <K> List<E> orderedMultiLoad( K[] naturalIds, SharedSessionContractImp
 
 		unorderedMultiLoad( naturalIds, session, lockOptions );
 
-		return handleResults( naturalIds, session, lockOptions );
+		return sortResults( naturalIds, session, lockOptions );
 	}
 
-	protected <K> List<E> handleResults( K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions ) {
-		List<E> results = new ArrayList<>(naturalIds.length);
+	protected <K> List<E> sortResults( K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions ) {
+		List<E> results = arrayList(naturalIds.length);
 		for ( int i = 0; i < naturalIds.length; i++ ) {
 			final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
 
 			Object id = persistenceContext.getNaturalIdResolutions().findCachedIdByNaturalId( naturalIds[i], getEntityDescriptor() );
 
 			// Id can be null if a non-existent natural id is requested
-			Object entity = id == null ? null
-					: persistenceContext.getEntity( new EntityKey( id, getEntityDescriptor().getEntityPersister() ) );
+			Object entity = (id == null) ? null : persistenceContext.getEntity( new EntityKey( id, getEntityDescriptor().getEntityPersister() ) );
 			if ( entity != null && !options.isReturnOfDeletedEntitiesEnabled() ) {
 				// make sure it is not DELETED
 				final EntityEntry entry = persistenceContext.getEntry( entity );
@@ -105,6 +118,38 @@ protected <K> List<E> handleResults( K[] naturalIds, SharedSessionContractImplem
 		return results;
 	}
 
+	private <K> Object[] checkPersistenceContextForCachedResults( K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions, List<E> results ) {
+		List<K> unresolvedIds = arrayList(naturalIds.length);
+
+		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
+		for ( int i = 0; i < naturalIds.length; i++ ) {
+
+			final Object normalizedNaturalId = getEntityDescriptor().getNaturalIdMapping().normalizeInput( naturalIds[i] );
+			Object id = persistenceContext.getNaturalIdResolutions().findCachedIdByNaturalId( normalizedNaturalId, getEntityDescriptor() );
+
+			// Id can be null if a non-existent natural id is requested, or a mutable natural id was changed and then deleted
+			Object entity = id == null ? null : persistenceContext.getEntity( new EntityKey( id, getEntityDescriptor().getEntityPersister() ) );
+
+			if ( entity != null ) {
+				// Entity is already in the persistence context
+				final EntityEntry entry = persistenceContext.getEntry( entity );
+				// either a managed entry, or a deleted one with returnDeleted enabled
+				if ( !entry.getStatus().isDeletedOrGone() || options.isReturnOfDeletedEntitiesEnabled() ) {
+					results.add( (E) persistenceContext.proxyFor(entity) );
+					upgradeLock( entity, entry, lockOptions, session.getSession().asEventSource() );
+				}
+			}
+			else {
+				// entity either doesn't exist or hasn't been loaded in the PC yet, in both cases we add the natural id
+				// to the ids that still need to be recovered, in case the id corresponds to a non-existent
+				// instance nothing will be in the results for it, which is ok in unordered multiload
+				unresolvedIds.add(naturalIds[i]);
+			}
+		}
+
+		return unresolvedIds.toArray( new Object[0] );
+	}
+
 	@Override
 	public EntityMappingType getLoadable() {
 		return getEntityDescriptor();

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiNaturalIdLoaderArrayParam.java

@@ -44,7 +44,7 @@ protected BasicAttributeMapping getNaturalIdAttribute()  {
 	}
 
 	@Override
-	public <K> List<E> unorderedMultiLoad( K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions ) {
+	public List<E> loadEntitiesWithUnresolvedIds(Object[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions) {
 
 		final SessionFactoryImplementor sessionFactory = session.getFactory();
 
@@ -77,17 +77,17 @@ public LockOptions getLockOptions() {
 				} );
 
 		return LoaderHelper.loadByArrayParameter(
-				naturalIds,
-				sqlAst,
-				jdbcSelectOperation,
-				jdbcParameter,
-				arrayJdbcMapping,
-				null,
-				null,
-				null,
-				lockOptions,
-				session.isDefaultReadOnly(),
-				session
+					naturalIds,
+					sqlAst,
+					jdbcSelectOperation,
+					jdbcParameter,
+					arrayJdbcMapping,
+					null,
+					null,
+					null,
+					lockOptions,
+					session.isDefaultReadOnly(),
+					session
 		);
 	}
 

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/MultiNaturalIdLoaderInPredicate.java

@@ -22,7 +22,7 @@ public MultiNaturalIdLoaderInPredicate(EntityMappingType entityDescriptor) {
 	}
 
 	@Override
-	public <K> List<E> unorderedMultiLoad(K[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions) {
+	public List<E> loadEntitiesWithUnresolvedIds(Object[] naturalIds, SharedSessionContractImplementor session, LockOptions lockOptions) {
 
 		final SessionFactoryImplementor sessionFactory = session.getFactory();