HHH-19734 Ensure shallow cached entity is not left uninitialized when bytecode enhanced

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

@@ -1102,6 +1102,10 @@ else if ( isResultInitializer() ) {
 				else if ( data.entityHolder.getEntityInitializer() != this ) {
 					data.setState( State.INITIALIZED );
 				}
+				else if ( data.shallowCached ) {
+					// For shallow cached entities, only the id is available, so ensure we load the data immediately
+					data.setInstance( data.entityInstanceForNotify = resolveEntityInstance( data ) );
+				}
 			}
 			else if ( ( entityFromExecutionContext = getEntityFromExecutionContext( data ) ) != null ) {
 				// This is the entity to refresh, so don't set the state to initialized
@@ -1214,7 +1218,7 @@ protected Object resolveEntityInstance(EntityInitializerData data) {
 			return resolved;
 		}
 		else {
-			if ( rowProcessingState.isQueryCacheHit() && entityDescriptor.useShallowQueryCacheLayout() ) {
+			if ( data.shallowCached ) {
 				// We must load the entity this way, because the query cache entry contains only the primary key
 				data.setState( State.INITIALIZED );
 				final SharedSessionContractImplementor session = rowProcessingState.getSession();

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/query/CachedQueryShallowWithDiscriminatorBytecodeEnhancedTest.java

@@ -0,0 +1,167 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.orm.test.jpa.query;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.Id;
+import jakarta.persistence.TypedQuery;
+import org.hibernate.annotations.Cache;
+import org.hibernate.annotations.CacheConcurrencyStrategy;
+import org.hibernate.annotations.CacheLayout;
+import org.hibernate.annotations.QueryCacheLayout;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.stat.Statistics;
+import org.hibernate.testing.bytecode.enhancement.extension.BytecodeEnhanced;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.JiraKey;
+import org.hibernate.testing.orm.junit.Jpa;
+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 static org.hibernate.jpa.HibernateHints.HINT_CACHEABLE;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@JiraKey("HHH-19734")
+@Jpa(
+		annotatedClasses = {
+				CachedQueryShallowWithDiscriminatorBytecodeEnhancedTest.Person.class
+		},
+		generateStatistics = true,
+		properties = {
+				@Setting(name = AvailableSettings.USE_QUERY_CACHE, value = "true"),
+				@Setting(name = AvailableSettings.USE_SECOND_LEVEL_CACHE, value = "true")
+		}
+)
+@BytecodeEnhanced
+public class CachedQueryShallowWithDiscriminatorBytecodeEnhancedTest {
+
+	public final static String HQL = "select p from Person p";
+
+	@BeforeEach
+	public void setUp(EntityManagerFactoryScope scope) {
+		scope.inTransaction(
+				em -> {
+					Person person = new Person( 1L );
+					person.setName( "Bob" );
+					em.persist( person );
+				}
+		);
+	}
+
+	@AfterEach
+	public void tearDown(EntityManagerFactoryScope scope) {
+		scope.inTransaction( em -> {
+			em.createQuery( "delete from Person" ).executeUpdate();
+		} );
+	}
+
+	@Test
+	public void testCacheableQuery(EntityManagerFactoryScope scope) {
+
+		Statistics stats = getStatistics( scope );
+		stats.clear();
+
+		// First time the query is executed, query and results are cached.
+		scope.inTransaction(
+				em -> {
+					loadPersons( em );
+
+					assertThatAnSQLQueryHasBeenExecuted( stats );
+
+					assertEquals( 0, stats.getQueryCacheHitCount() );
+					assertEquals( 1, stats.getQueryCacheMissCount() );
+					assertEquals( 1, stats.getQueryCachePutCount() );
+
+					assertEquals( 0, stats.getSecondLevelCacheHitCount() );
+					assertEquals( 0, stats.getSecondLevelCacheMissCount() );
+					assertEquals( 0, stats.getSecondLevelCachePutCount() );
+				}
+		);
+
+		stats.clear();
+
+		// Second time the query is executed, list of entities are read from query cache
+
+		scope.inTransaction(
+				em -> {
+					// Create a person proxy in the persistence context to trigger the HHH-19734 error
+					em.getReference( Person.class, 1L );
+
+					loadPersons( em );
+
+					assertThatNoSQLQueryHasBeenExecuted( stats );
+
+					assertEquals( 1, stats.getQueryCacheHitCount() );
+					assertEquals( 0, stats.getQueryCacheMissCount() );
+					assertEquals( 0, stats.getQueryCachePutCount() );
+
+					assertEquals( 1, stats.getSecondLevelCacheHitCount() );
+					assertEquals( 0, stats.getSecondLevelCacheMissCount() );
+					assertEquals( 0, stats.getSecondLevelCachePutCount() );
+				}
+		);
+
+	}
+
+	private static Statistics getStatistics(EntityManagerFactoryScope scope) {
+		return ((SessionFactoryImplementor) scope.getEntityManagerFactory()).getStatistics();
+	}
+
+	private static void loadPersons(EntityManager em) {
+		TypedQuery<Person> query = em.createQuery( HQL, Person.class )
+				.setHint( HINT_CACHEABLE, true );
+		Person person = query.getSingleResult();
+		assertEquals( 1L, person.getId() );
+		assertEquals( "Bob", person.getName() );
+	}
+
+	private static void assertThatAnSQLQueryHasBeenExecuted(Statistics stats) {
+		assertEquals( 1, stats.getQueryStatistics( HQL ).getExecutionCount() );
+	}
+
+	private static void assertThatNoSQLQueryHasBeenExecuted(Statistics stats) {
+		assertEquals( 0, stats.getQueryStatistics( HQL ).getExecutionCount() );
+	}
+
+	@Entity(name = "Person")
+	@Cache(usage = CacheConcurrencyStrategy.READ_WRITE)
+	@QueryCacheLayout(layout = CacheLayout.SHALLOW)
+	public static class Person {
+		@Id
+		private Long id;
+		private String name;
+
+		public Person() {
+			super();
+		}
+
+		public Person(Long id) {
+			this.id = id;
+		}
+
+		public Long getId() {
+			return id;
+		}
+
+		public void setId(Long id) {
+			this.id = id;
+		}
+
+		public String getName() {
+			return name;
+		}
+
+		public void setName(String name) {
+			this.name = name;
+		}
+	}
+
+}

hibernate-testing/src/main/java/org/hibernate/testing/orm/jpa/PersistenceUnitInfoImpl.java

@@ -40,6 +40,7 @@ public class PersistenceUnitInfoImpl implements PersistenceUnitInfo {
 	private List<String> mappingFiles;
 	private List<String> managedClassNames;
 	private boolean excludeUnlistedClasses;
+	private ClassLoader classLoader;
 
 	public PersistenceUnitInfoImpl(String name) {
 		this.name = name;
@@ -120,6 +121,15 @@ public void setExcludeUnlistedClasses(boolean excludeUnlistedClasses) {
 		this.excludeUnlistedClasses = excludeUnlistedClasses;
 	}
 
+	@Override
+	public ClassLoader getClassLoader() {
+		return classLoader;
+	}
+
+	public void setClassLoader(ClassLoader classLoader) {
+		this.classLoader = classLoader;
+	}
+
 	@Override
 	public String getPersistenceXMLSchemaVersion() {
 		return null;
@@ -145,11 +155,6 @@ public URL getPersistenceUnitRootUrl() {
 		return null;
 	}
 
-	@Override
-	public ClassLoader getClassLoader() {
-		return null;
-	}
-
 	@Override
 	public void addTransformer(ClassTransformer transformer) {
 

hibernate-testing/src/main/java/org/hibernate/testing/orm/junit/EntityManagerFactoryExtension.java

@@ -90,6 +90,9 @@ public static EntityManagerFactoryScope findEntityManagerFactoryScope(
 						pui.getProperties().put( key, value )
 		);
 
+		// Use the context class loader for entity loading if configured,
+		// to make enhancement work for tests
+		pui.setClassLoader( Thread.currentThread().getContextClassLoader() );
 		pui.setTransactionType( emfAnn.transactionType() );
 		pui.setCacheMode( emfAnn.sharedCacheMode() );
 		pui.setValidationMode( emfAnn.validationMode() );