HHH-18689 Maintain proxy targets when converting cache entries

hibernate-core/src/main/java/org/hibernate/engine/internal/TwoPhaseLoad.java

@@ -5,10 +5,16 @@
 package org.hibernate.engine.internal;
 
 import org.hibernate.LockMode;
+import org.hibernate.engine.spi.EntityEntry;
+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.engine.spi.Status;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.proxy.LazyInitializer;
+
+import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
 
 /**
  * Functionality relating to the Hibernate two-phase loading process, that may be reused by persisters
@@ -37,16 +43,27 @@ public static void addUninitializedCachedEntity(
 			final LockMode lockMode,
 			final Object version,
 			final SharedSessionContractImplementor session) {
-		session.getPersistenceContextInternal().addEntity(
+		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
+		final EntityHolder entityHolder = persistenceContext.addEntityHolder( key, object );
+		final EntityEntry entityEntry = persistenceContext.addEntry(
 				object,
 				Status.LOADING,
 				null,
-				key,
+				null,
+				key.getIdentifier(),
 				version,
 				lockMode,
 				true,
 				persister,
 				false
 		);
+		entityHolder.setEntityEntry( entityEntry );
+		final Object proxy = entityHolder.getProxy();
+		if ( proxy != null ) {
+			// there is already a proxy for this impl
+			final LazyInitializer lazyInitializer = extractLazyInitializer( proxy );
+			assert lazyInitializer != null;
+			lazyInitializer.setImplementation( object );
+		}
 	}
 }

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

@@ -31,6 +31,7 @@
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.pretty.MessageHelper;
+import org.hibernate.proxy.LazyInitializer;
 import org.hibernate.sql.results.LoadingLogger;
 import org.hibernate.stat.internal.StatsHelper;
 import org.hibernate.stat.spi.StatisticsImplementor;
@@ -377,16 +378,17 @@ private Object convertCacheEntryToEntity(
 		final EntityPersister subclassPersister =
 				factory.getRuntimeMetamodels().getMappingMetamodel()
 						.getEntityDescriptor( entry.getSubclass() );
+		final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
+		final EntityHolder oldHolder = persistenceContext.getEntityHolder( entityKey );
 
 		final Object entity;
 		if ( instanceToLoad != null ) {
 			entity = instanceToLoad;
 		}
 		else {
-			final EntityHolder holder = source.getPersistenceContextInternal().getEntityHolder( entityKey );
-			if ( holder != null && holder.getEntity() != null ) {
+			if ( oldHolder != null && oldHolder.getEntity() != null ) {
 				// Use the entity which might already be
-				entity = holder.getEntity();
+				entity = oldHolder.getEntity();
 			}
 			else {
 				entity = source.instantiate( subclassPersister, entityId );
@@ -407,13 +409,34 @@ private Object convertCacheEntryToEntity(
 		}
 
 		// make it circular-reference safe
-		TwoPhaseLoad.addUninitializedCachedEntity(
-				entityKey,
-				entity,
-				subclassPersister,
-				LockMode.NONE,
-				entry.getVersion(),
-				source
+		final EntityHolder holder = persistenceContext.addEntityHolder( entityKey, entity );
+		final Object proxy = holder.getProxy();
+		final boolean isReadOnly;
+		if ( proxy != null ) {
+			// there is already a proxy for this impl
+			// only set the status to read-only if the proxy is read-only
+			final LazyInitializer lazyInitializer = extractLazyInitializer( proxy );
+			assert lazyInitializer != null;
+			lazyInitializer.setImplementation( entity );
+
+			isReadOnly = lazyInitializer.isReadOnly();
+		}
+		else {
+			isReadOnly = source.isDefaultReadOnly();
+		}
+		holder.setEntityEntry(
+				persistenceContext.addEntry(
+						entity,
+						Status.LOADING,
+						null,
+						null,
+						entityKey.getIdentifier(),
+						entry.getVersion(),
+						LockMode.NONE,
+						true,
+						persister,
+						false
+				)
 		);
 
 		final Type[] types = subclassPersister.getPropertyTypes();
@@ -438,32 +461,20 @@ private Object convertCacheEntryToEntity(
 		final Object version = getVersion( values, subclassPersister );
 		LOG.tracef( "Cached Version : %s", version );
 
-		final PersistenceContext persistenceContext = source.getPersistenceContext();
-
-		final Object proxy = persistenceContext.getProxy( entityKey );
-		final boolean isReadOnly;
-		if ( proxy != null ) {
-			// there is already a proxy for this impl
-			// only set the status to read-only if the proxy is read-only
-			isReadOnly = extractLazyInitializer( proxy ).isReadOnly();
-		}
-		else {
-			isReadOnly = source.isDefaultReadOnly();
-		}
-
-		final EntityEntry entityEntry = persistenceContext.addEntry(
-				entity,
-				isReadOnly ? Status.READ_ONLY : Status.MANAGED,
-				values,
-				null,
-				entityId,
-				version,
-				LockMode.NONE,
-				true,
-				subclassPersister,
-				false
+		holder.setEntityEntry(
+				persistenceContext.addEntry(
+						entity,
+						isReadOnly ? Status.READ_ONLY : Status.MANAGED,
+						values,
+						null,
+						entityId,
+						version,
+						LockMode.NONE,
+						true,
+						subclassPersister,
+						false
+				)
 		);
-		persistenceContext.getEntityHolder( entityKey ).setEntityEntry( entityEntry );
 		subclassPersister.afterInitialize( entity, source );
 		persistenceContext.initializeNonLazyCollections();
 

hibernate-core/src/test/java/org/hibernate/orm/test/cache/QueryCacheIncompleteTest.java

@@ -0,0 +1,114 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.cache;
+
+import org.hibernate.CacheMode;
+import org.hibernate.Session;
+import org.hibernate.annotations.Cache;
+import org.hibernate.annotations.CacheConcurrencyStrategy;
+import org.hibernate.cfg.AvailableSettings;
+
+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.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@DomainModel(annotatedClasses = {
+		QueryCacheIncompleteTest.Admin.class,
+})
+@SessionFactory
+@ServiceRegistry(
+		settings = {
+				@Setting(name = AvailableSettings.USE_SECOND_LEVEL_CACHE, value = "true"),
+				@Setting(name = AvailableSettings.USE_QUERY_CACHE, value = "true"),
+				@Setting(name = AvailableSettings.QUERY_CACHE_LAYOUT, value = "FULL")
+		}
+)
+@JiraKey(value = "HHH-18689")
+public class QueryCacheIncompleteTest {
+
+	private Long adminId;
+
+	@BeforeAll
+	public void setUp(SessionFactoryScope scope) {
+		adminId = scope.fromTransaction(
+				session -> {
+					Admin admin = new Admin();
+					admin.setAge( 42 );
+					session.persist( admin );
+					return admin.getId();
+				}
+		);
+	}
+
+	@Test
+	void testQueryWithEmbeddableParameter(SessionFactoryScope scope) {
+		scope.inTransaction(
+				session -> {
+					// load uninitialized proxy
+					session.getReference( Admin.class, adminId );
+					// load entity
+					var multiLoader = session.byMultipleIds( Admin.class );
+					multiLoader.with( CacheMode.NORMAL );
+					multiLoader.multiLoad( adminId );
+
+					// store in query cache
+					Admin admin = queryAdmin( session );
+					assertThat( admin.getAge() ).isEqualTo( 42 );
+				}
+		);
+
+		scope.inTransaction(
+				session -> {
+					// use query cache
+					Admin admin = queryAdmin( session );
+					assertThat( admin.getAge() ).isEqualTo( 42 );
+				}
+		);
+	}
+
+	private Admin queryAdmin(Session s) {
+		return s.createQuery( "from Admin", Admin.class ).setCacheable( true ).getSingleResult();
+	}
+
+	@Entity(name = "Admin")
+	@Cache(usage = CacheConcurrencyStrategy.READ_WRITE)
+	public static class Admin {
+
+		@Id
+		@GeneratedValue
+		private Long id;
+
+		@Column(nullable = false)
+		private int age;
+
+		public Long getId() {
+			return id;
+		}
+
+		public void setId(Long id) {
+			this.id = id;
+		}
+
+		public int getAge() {
+			return age;
+		}
+
+		public void setAge(int age) {
+			this.age = age;
+		}
+	}
+}