HHH-18131 Composite identifiers with associations stopped working with @IdClass

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultMergeEventListener.java

@@ -162,7 +162,7 @@ public class DefaultMergeEventListener
 		final Object originalId;
 		if ( entry == null ) {
 			final EntityPersister persister = source.getEntityPersister( event.getEntityName(), entity );
-			originalId = persister.getIdentifier( entity, source );
+			originalId = persister.getIdentifier( entity, copiedAlready );
 			if ( originalId != null ) {
 				final EntityKey entityKey;
 				if ( persister.getIdentifierType() instanceof ComponentType ) {

hibernate-core/src/main/java/org/hibernate/event/spi/MergeContext.java

@@ -364,4 +364,8 @@ public class MergeContext implements Map<Object,Object> {
 		// Entity was not found in current persistence context. Use Object#toString() method.
 		return "[" + entity + "]";
 	}
+
+	public EventSource getEventSource() {
+		return session;
+	}
 }

hibernate-core/src/main/java/org/hibernate/internal/util/EntityPrinter.java

@@ -51,7 +51,7 @@ public final class EntityPrinter {
 			result.put(
 					entityPersister.getIdentifierPropertyName(),
 					entityPersister.getIdentifierType().toLoggableString(
-							entityPersister.getIdentifier( entity, null ),
+							entityPersister.getIdentifier( entity ),
 							factory
 					)
 			);

hibernate-core/src/main/java/org/hibernate/jpa/internal/PersistenceUnitUtilImpl.java

@@ -159,7 +159,7 @@ public class PersistenceUnitUtilImpl implements PersistenceUnitUtil, Serializabl
 		catch (MappingException ex) {
 			throw new IllegalArgumentException( entityClass.getName() + " is not an entity", ex );
 		}
-		return persister.getIdentifier( entity, null );
+		return persister.getIdentifier( entity );
 	}
 
 	private Object getVersionFromPersister(Object entity) {

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/EntityIdentifierMapping.java

@@ -8,6 +8,7 @@ import org.hibernate.TransientObjectException;
 import org.hibernate.engine.internal.ForeignKeys;
 import org.hibernate.engine.spi.IdentifierValue;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 
 import jakarta.persistence.EmbeddedId;
 import jakarta.persistence.Id;
@@ -66,6 +67,15 @@ public interface EntityIdentifierMapping extends ValuedModelPart {
 	 */
 	Object getIdentifier(Object entity);
 
+	/**
+	 * Extract the identifier from an instance of the entity
+	 *
+	 * It's supposed to be use during the merging process
+	 */
+	default Object getIdentifier(Object entity, MergeContext mergeContext){
+		return getIdentifier( entity );
+	}
+
 	/**
 	 * Return the identifier of the persistent or transient object, or throw
 	 * an exception if the instance is "unsaved"

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/InverseNonAggregatedIdentifierMapping.java

@@ -12,6 +12,7 @@ import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.metamodel.mapping.AttributeMapping;
 import org.hibernate.metamodel.mapping.EmbeddableMappingType;
@@ -194,6 +195,11 @@ public class InverseNonAggregatedIdentifierMapping extends EmbeddedAttributeMapp
 
 	@Override
 	public Object getIdentifier(Object entity) {
+		return getIdentifier( entity, null );
+	}
+
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
 		if ( hasContainingClass() ) {
 			final Object id = identifierValueMapper.getRepresentationStrategy().getInstantiator().instantiate(
 					null,
@@ -214,18 +220,18 @@ public class InverseNonAggregatedIdentifierMapping extends EmbeddedAttributeMapp
 					}
 				}
 				//JPA 2 @MapsId + @IdClass points to the pk of the entity
-				else if ( attributeMapping instanceof ToOneAttributeMapping
+				else if ( attributeMapping instanceof ToOneAttributeMapping toOneAttributeMapping
 						&& !( identifierValueMapper.getAttributeMapping( i ) instanceof ToOneAttributeMapping ) ) {
-					final ToOneAttributeMapping toOneAttributeMapping = (ToOneAttributeMapping) attributeMapping;
+					final Object toOne = getIfMerged( o, mergeContext );
 					final ModelPart targetPart = toOneAttributeMapping.getForeignKeyDescriptor().getPart(
 							toOneAttributeMapping.getSideNature().inverse()
 					);
 					if ( targetPart.isEntityIdentifierMapping() ) {
-						propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( o );
+						propertyValues[i] = ( (EntityIdentifierMapping) targetPart )
+								.getIdentifier( toOne, mergeContext );
 					}
 					else {
-						propertyValues[i] = o;
+						propertyValues[i] = toOne;
-						assert false;
 					}
 				}
 				else {
@@ -240,6 +246,16 @@ public class InverseNonAggregatedIdentifierMapping extends EmbeddedAttributeMapp
 		}
 	}
 
+	private static Object getIfMerged(Object o, MergeContext mergeContext) {
+		if ( mergeContext != null ) {
+			final Object merged = mergeContext.get( o );
+			if ( merged != null ) {
+				return merged;
+			}
+		}
+		return o;
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		final Object[] propertyValues = new Object[identifierValueMapper.getNumberOfAttributeMappings()];

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/NonAggregatedIdentifierMappingImpl.java

@@ -13,6 +13,7 @@ import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.mapping.Component;
 import org.hibernate.mapping.RootClass;
@@ -235,6 +236,11 @@ public class NonAggregatedIdentifierMappingImpl extends AbstractCompositeIdentif
 
 	@Override
 	public Object getIdentifier(Object entity) {
+		return getIdentifier( entity, null );
+	}
+
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
 		if ( hasContainingClass() ) {
 			final LazyInitializer lazyInitializer = HibernateProxy.extractLazyInitializer( entity );
 			if ( lazyInitializer != null ) {
@@ -255,18 +261,17 @@ public class NonAggregatedIdentifierMappingImpl extends AbstractCompositeIdentif
 					}
 				}
 				//JPA 2 @MapsId + @IdClass points to the pk of the entity
-				else if ( attributeMapping instanceof ToOneAttributeMapping
+				else if ( attributeMapping instanceof ToOneAttributeMapping toOneAttributeMapping
 						&& !( identifierValueMapper.getAttributeMapping( i ) instanceof ToOneAttributeMapping ) ) {
-					final ToOneAttributeMapping toOneAttributeMapping = (ToOneAttributeMapping) attributeMapping;
+					final Object toOne = getIfMerged( o, mergeContext );
 					final ModelPart targetPart = toOneAttributeMapping.getForeignKeyDescriptor().getPart(
 							toOneAttributeMapping.getSideNature().inverse()
 					);
 					if ( targetPart.isEntityIdentifierMapping() ) {
-						propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( o );
+						propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( toOne, mergeContext );
 					}
 					else {
-						propertyValues[i] = o;
+						propertyValues[i] = toOne;
-						assert false;
 					}
 				}
 				else {
@@ -283,6 +288,16 @@ public class NonAggregatedIdentifierMappingImpl extends AbstractCompositeIdentif
 		}
 	}
 
+	private static Object getIfMerged(Object o, MergeContext mergeContext) {
+		if ( mergeContext != null ) {
+			final Object merged = mergeContext.get( o );
+			if ( merged != null ) {
+				return merged;
+			}
+		}
+		return o;
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		final Object[] propertyValues = new Object[identifierValueMapper.getNumberOfAttributeMappings()];

hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java

@@ -94,6 +94,7 @@ import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.event.spi.LoadEvent;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.generator.BeforeExecutionGenerator;
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.Generator;
@@ -4297,6 +4298,11 @@ public abstract class AbstractEntityPersister
 		return identifierMapping.getIdentifier( entity );
 	}
 
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
+		return identifierMapping.getIdentifier( entity, mergeContext );
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		identifierMapping.setIdentifier( entity, id, session );

hibernate-core/src/main/java/org/hibernate/persister/entity/EntityPersister.java

@@ -28,6 +28,7 @@ import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.event.spi.EventSource;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.generator.BeforeExecutionGenerator;
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.Generator;
@@ -1129,6 +1130,24 @@ public interface EntityPersister extends EntityMappingType, EntityMutationTarget
 	 */
 	Object getIdentifier(Object entity, SharedSessionContractImplementor session);
 
+	/**
+	 * Get the identifier of an instance from the object's identifier property.
+	 * Throw an exception if it has no identifier property.
+	 *
+	 * It's supposed to be use during the merging process
+	 */
+	default Object getIdentifier(Object entity, MergeContext mergeContext) {
+		return getIdentifier( entity, mergeContext.getEventSource() );
+	}
+
+	/**
+	 * Get the identifier of an instance from the object's identifier property.
+	 * Throw an exception if it has no identifier property.
+	 */
+	default Object getIdentifier(Object entity) {
+		return getIdentifier( entity, (SharedSessionContractImplementor) null );
+	}
+
 	/**
 	 * Inject the identifier value into the given entity.
 	 */

hibernate-core/src/main/java/org/hibernate/query/derived/AnonymousTupleBasicEntityIdentifierMapping.java

@@ -7,6 +7,7 @@ package org.hibernate.query.derived;
 import org.hibernate.Incubating;
 import org.hibernate.engine.spi.IdentifierValue;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.metamodel.mapping.BasicEntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.JdbcMapping;
 import org.hibernate.metamodel.mapping.ManagedMappingType;
@@ -49,6 +50,11 @@ public class AnonymousTupleBasicEntityIdentifierMapping
 		return delegate.getIdentifier( entity );
 	}
 
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
+		return delegate.getIdentifier( entity, mergeContext );
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		delegate.setIdentifier( entity, id, session );

hibernate-core/src/main/java/org/hibernate/query/derived/AnonymousTupleEmbeddedEntityIdentifierMapping.java

@@ -10,6 +10,7 @@ import java.util.Set;
 import org.hibernate.Incubating;
 import org.hibernate.engine.spi.IdentifierValue;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.metamodel.mapping.CompositeIdentifierMapping;
 import org.hibernate.metamodel.mapping.EmbeddableMappingType;
 import org.hibernate.metamodel.mapping.internal.SingleAttributeIdentifierMapping;
@@ -68,6 +69,11 @@ public class AnonymousTupleEmbeddedEntityIdentifierMapping extends AnonymousTupl
 		return delegate.getIdentifier( entity );
 	}
 
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
+		return delegate.getIdentifier( entity, mergeContext );
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		delegate.setIdentifier( entity, id, session );

hibernate-core/src/main/java/org/hibernate/query/derived/AnonymousTupleNonAggregatedEntityIdentifierMapping.java

@@ -12,6 +12,7 @@ import org.hibernate.engine.FetchStyle;
 import org.hibernate.engine.FetchTiming;
 import org.hibernate.engine.spi.IdentifierValue;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.MergeContext;
 import org.hibernate.metamodel.mapping.EmbeddableMappingType;
 import org.hibernate.metamodel.mapping.NonAggregatedIdentifierMapping;
 import org.hibernate.metamodel.mapping.internal.IdClassEmbeddable;
@@ -77,6 +78,12 @@ public class AnonymousTupleNonAggregatedEntityIdentifierMapping extends Anonymou
 		return delegate.getIdentifier( entity );
 	}
 
+
+	@Override
+	public Object getIdentifier(Object entity, MergeContext mergeContext) {
+		return delegate.getIdentifier( entity, mergeContext );
+	}
+
 	@Override
 	public void setIdentifier(Object entity, Object id, SharedSessionContractImplementor session) {
 		delegate.setIdentifier( entity, id, session );

hibernate-core/src/main/java/org/hibernate/type/AnyType.java

@@ -146,7 +146,7 @@ public class AnyType extends AbstractType implements CompositeType, AssociationT
 		final EntityPersister concretePersister = guessEntityPersister( entity, factory );
 		return concretePersister == null
 				? null
-				: concretePersister.getIdentifier( entity, null );
+				: concretePersister.getIdentifier( entity );
 	}
 
 	private EntityPersister guessEntityPersister(Object object, SessionFactoryImplementor factory) {

hibernate-core/src/main/java/org/hibernate/type/EntityType.java

@@ -281,7 +281,7 @@ public abstract class EntityType extends AbstractType implements AssociationType
 		final EntityPersister concretePersister = getAssociatedEntityPersister( factory );
 		return concretePersister == null
 				? null
-				: concretePersister.getIdentifier( entity, null );
+				: concretePersister.getIdentifier( entity );
 	}
 
 	@Override
@@ -350,7 +350,7 @@ public abstract class EntityType extends AbstractType implements AssociationType
 			else {
 				final Class<?> mappedClass = persister.getMappedClass();
 				if ( mappedClass.isAssignableFrom( x.getClass() ) ) {
-					id = persister.getIdentifier( x, null );
+					id = persister.getIdentifier( x );
 				}
 				else {
 					id = x;
@@ -384,7 +384,7 @@ public abstract class EntityType extends AbstractType implements AssociationType
 		}
 		else {
 			if ( mappedClass.isAssignableFrom( x.getClass() ) ) {
-				xid = persister.getIdentifier( x, null );
+				xid = persister.getIdentifier( x );
 			}
 			else {
 				//JPA 2 case where @IdClass contains the id and not the associated entity
@@ -399,7 +399,7 @@ public abstract class EntityType extends AbstractType implements AssociationType
 		}
 		else {
 			if ( mappedClass.isAssignableFrom( y.getClass() ) ) {
-				yid = persister.getIdentifier( y, null );
+				yid = persister.getIdentifier( y );
 			}
 			else {
 				//JPA 2 case where @IdClass contains the id and not the associated entity
@@ -561,7 +561,7 @@ public abstract class EntityType extends AbstractType implements AssociationType
 				id = lazyInitializer.getInternalIdentifier();
 			}
 			else {
-				id = persister.getIdentifier( value, null );
+				id = persister.getIdentifier( value );
 			}
 
 			result.append( '#' )