HHH-18645 Handle proxies when resolving from existing entity in batch initializer

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/AbstractBatchEntitySelectFetchInitializer.java

@@ -8,6 +8,7 @@
 import org.hibernate.FetchNotFoundException;
 import org.hibernate.Hibernate;
 import org.hibernate.annotations.NotFoundAction;
+import org.hibernate.bytecode.enhance.spi.interceptor.EnhancementAsProxyLazinessInterceptor;
 import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
@@ -29,7 +30,9 @@
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import static org.hibernate.engine.internal.ManagedTypeHelper.isPersistentAttributeInterceptable;
 import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
+import static org.hibernate.sql.results.graph.entity.internal.EntityInitializerImpl.getAttributeInterceptor;
 
 public abstract class AbstractBatchEntitySelectFetchInitializer<Data extends AbstractBatchEntitySelectFetchInitializer.AbstractBatchEntitySelectFetchInitializerData>
 		extends EntitySelectFetchInitializer<Data> implements EntityInitializer<Data> {
@@ -109,6 +112,10 @@ public void resolveInstance(Data data) {
 				return;
 			}
 		}
+		resolveInstanceFromIdentifier( data );
+	}
+
+	protected void resolveInstanceFromIdentifier(Data data) {
 		if ( data.batchDisabled ) {
 			initialize( data );
 		}
@@ -135,21 +142,34 @@ public void resolveInstance(Object instance, Data data) {
 		// Only need to extract the identifier if the identifier has a many to one
 		final LazyInitializer lazyInitializer = extractLazyInitializer( instance );
 		data.entityKey = null;
+		data.entityIdentifier = null;
 		if ( lazyInitializer == null ) {
-			// Entity is initialized
-			data.setState( State.INITIALIZED );
-			if ( keyIsEager ) {
+			// Entity is most probably initialized
+			data.setInstance( instance );
+			if ( concreteDescriptor.getBytecodeEnhancementMetadata().isEnhancedForLazyLoading()
+					&& isPersistentAttributeInterceptable( instance )
+					&& getAttributeInterceptor( instance ) instanceof EnhancementAsProxyLazinessInterceptor enhancementInterceptor ) {
+				if ( enhancementInterceptor.isInitialized() ) {
+					data.setState( State.INITIALIZED );
+				}
+				else {
+					data.setState( State.RESOLVED );
+					data.entityIdentifier = enhancementInterceptor.getIdentifier();
+				}
+			}
+			else {
+				// If the entity initializer is null, we know the entity is fully initialized,
+				// otherwise it will be initialized by some other initializer
+				data.setState( State.RESOLVED );
+				data.entityIdentifier = concreteDescriptor.getIdentifier( instance, rowProcessingState.getSession() );
+			}
+			if ( keyIsEager && data.entityIdentifier == null ) {
 				data.entityIdentifier = concreteDescriptor.getIdentifier( instance, rowProcessingState.getSession() );
 			}
-			data.setInstance( instance );
 		}
 		else if ( lazyInitializer.isUninitialized() ) {
 			data.setState( State.RESOLVED );
-			if ( keyIsEager ) {
-				data.entityIdentifier = lazyInitializer.getIdentifier();
-			}
-			// Resolve and potentially create the entity instance
-			registerToBatchFetchQueue( data );
+			data.entityIdentifier = lazyInitializer.getIdentifier();
 		}
 		else {
 			// Entity is initialized
@@ -159,6 +179,10 @@ else if ( lazyInitializer.isUninitialized() ) {
 			}
 			data.setInstance( lazyInitializer.getImplementation() );
 		}
+
+		if ( data.getState() == State.RESOLVED ) {
+			resolveInstanceFromIdentifier( data );
+		}
 		if ( keyIsEager ) {
 			final Initializer<?> initializer = keyAssembler.getInitializer();
 			assert initializer != null;

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/BatchEntitySelectFetchInitializer.java

@@ -64,15 +64,17 @@ protected InitializerData createInitializerData(RowProcessingState rowProcessing
 	@Override
 	protected void registerResolutionListener(BatchEntitySelectFetchInitializerData data) {
 		final RowProcessingState rowProcessingState = data.getRowProcessingState();
-		final InitializerData owningData = owningEntityInitializer.getData( rowProcessingState );
+		final InitializerData owningData = owningEntityInitializer.getData( rowProcessingState );HashMap<EntityKey, List<ParentInfo>> toBatchLoad = data.toBatchLoad;
+		if ( toBatchLoad == null ) {
+			toBatchLoad = data.toBatchLoad = new HashMap<>();
+		}
+		// Always register the entity key for resolution
+		final List<ParentInfo> parentInfos = toBatchLoad.computeIfAbsent( data.entityKey, key -> new ArrayList<>() );
 		final AttributeMapping parentAttribute;
+		// But only add the parent info if the parent entity is not already initialized
 		if ( owningData.getState() != State.INITIALIZED
 				&& ( parentAttribute = parentAttributes[owningEntityInitializer.getConcreteDescriptor( owningData ).getSubclassId()] ) != null ) {
-			HashMap<EntityKey, List<ParentInfo>> toBatchLoad = data.toBatchLoad;
-			if ( toBatchLoad == null ) {
-				toBatchLoad = data.toBatchLoad = new HashMap<>();
-			}
-			toBatchLoad.computeIfAbsent( data.entityKey, key -> new ArrayList<>() ).add(
+			parentInfos.add(
 					new ParentInfo(
 							owningEntityInitializer.getTargetInstance( owningData ),
 							parentAttribute.getStateArrayPosition()

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/EntityInitializerImpl.java

@@ -1797,7 +1797,7 @@ protected void forEachSubInitializer(BiConsumer<Initializer<?>, RowProcessingSta
 		}
 	}
 
-	private static PersistentAttributeInterceptor getAttributeInterceptor(Object entity) {
+	public static PersistentAttributeInterceptor getAttributeInterceptor(Object entity) {
 		return asPersistentAttributeInterceptable( entity ).$$_hibernate_getInterceptor();
 	}
 

hibernate-core/src/test/java/org/hibernate/orm/test/bytecode/enhancement/batch/BatchLazyProxyTest.java

@@ -0,0 +1,207 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.bytecode.enhancement.batch;
+
+import jakarta.persistence.CascadeType;
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.FetchType;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+import jakarta.persistence.JoinColumn;
+import jakarta.persistence.ManyToOne;
+import jakarta.persistence.OneToMany;
+import jakarta.persistence.OneToOne;
+import jakarta.persistence.Table;
+import org.hibernate.Hibernate;
+import org.hibernate.annotations.BatchSize;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.graph.GraphSemantic;
+import org.hibernate.testing.bytecode.enhancement.extension.BytecodeEnhanced;
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.JiraKey;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+@DomainModel(
+		annotatedClasses = {
+				BatchLazyProxyTest.User.class,
+				BatchLazyProxyTest.UserInfo.class,
+				BatchLazyProxyTest.Phone.class,
+		}
+
+)
+@SessionFactory
+@ServiceRegistry(
+		settings = {
+				@Setting(name = AvailableSettings.DEFAULT_BATCH_FETCH_SIZE, value = "100")
+
+		}
+)
+@JiraKey("HHH-18645")
+@BytecodeEnhanced(runNotEnhancedAsWell = true)
+public class BatchLazyProxyTest {
+
+	@BeforeEach
+	public void setUp(SessionFactoryScope scope) {
+		scope.inTransaction(
+				session -> {
+					UserInfo info = new UserInfo( "info" );
+					Phone phone = new Phone( "123456" );
+					info.addPhone( phone );
+					User user = new User( 1L, "user1", info );
+					session.persist( user );
+				}
+		);
+	}
+
+	@AfterEach
+	public void tearDown(SessionFactoryScope scope) {
+		scope.inTransaction(
+				session -> {
+					session.createMutationQuery( "delete User" ).executeUpdate();
+					session.createMutationQuery( "delete Phone" ).executeUpdate();
+					session.createMutationQuery( "delete UserInfo" ).executeUpdate();
+				}
+		);
+	}
+
+	@Test
+	public void testBatchInitialize(SessionFactoryScope scope) {
+		scope.inTransaction(
+				session -> {
+					User user = session.createQuery( "select u from User u where u.id = :id", User.class )
+							.setEntityGraph( session.createEntityGraph( User.class ), GraphSemantic.FETCH )
+							.setParameter( "id", 1L )
+							.getSingleResult();
+					assertThat( Hibernate.isInitialized( user.getInfo() ) ).isFalse();
+					session.createQuery( "select u from User u where u.id = :id", User.class )
+							.setParameter( "id", 1L )
+							.getSingleResult();
+					assertThat( Hibernate.isInitialized( user.getInfo() ) ).isTrue();
+				}
+		);
+	}
+
+	@Entity(name = "User")
+	@Table(name = "USER_TABLE")
+	@BatchSize(size = 5)
+	public static class User {
+
+		@Id
+		private Long id;
+
+		@Column
+		private String name;
+
+		@OneToOne(cascade = CascadeType.ALL, fetch = FetchType.EAGER)
+		@JoinColumn(name = "INFO_ID", referencedColumnName = "ID")
+		private UserInfo info;
+
+		public User() {
+		}
+
+		public User(long id, String name, UserInfo info) {
+			this.id = id;
+			this.name = name;
+			this.info = info;
+			info.user = this;
+		}
+
+		public long getId() {
+			return id;
+		}
+
+		public String getName() {
+			return name;
+		}
+
+		public UserInfo getInfo() {
+			return info;
+		}
+	}
+
+	@Entity(name = "UserInfo")
+	public static class UserInfo {
+		@Id
+		@GeneratedValue
+		private Long id;
+
+		@OneToOne(mappedBy = "info", fetch = FetchType.LAZY)
+		private User user;
+
+		private String info;
+
+		@OneToMany(mappedBy = "info", cascade = CascadeType.ALL, fetch = FetchType.LAZY)
+		private List<Phone> phoneList;
+
+		public long getId() {
+			return id;
+		}
+
+		public UserInfo() {
+		}
+
+		public UserInfo(String info) {
+			this.info = info;
+		}
+
+		public User getUser() {
+			return user;
+		}
+
+		public String getInfo() {
+			return info;
+		}
+
+		public List<Phone> getPhoneList() {
+			return phoneList;
+		}
+
+		public void addPhone(Phone phone) {
+			if ( phoneList == null ) {
+				phoneList = new ArrayList<>();
+			}
+			this.phoneList.add( phone );
+			phone.info = this;
+		}
+	}
+
+	@Entity(name = "Phone")
+	public static class Phone {
+		@Id
+		@Column(name = "PHONE_NUMBER")
+		private String number;
+
+		@ManyToOne
+		@JoinColumn(name = "INFO_ID")
+		private UserInfo info;
+
+		public Phone() {
+		}
+
+		public Phone(String number) {
+			this.number = number;
+		}
+
+		public String getNumber() {
+			return number;
+		}
+
+		public UserInfo getInfo() {
+			return info;
+		}
+	}
+}