HHH-19096 Adjust SelectionQuery#setEntityGraph(..) to accept entity graphs of supertypes

Author: marko-bekhta

hibernate-core/src/main/java/org/hibernate/Session.java

@@ -1311,7 +1311,8 @@ public interface Session extends SharedSessionContract, EntityManager {
 	 * @apiNote This method returns {@code RootGraph<?>}, requiring an
 	 * unchecked typecast before use. It's cleaner to obtain a graph using
 	 * the static metamodel for the class which defines the graph, or by
-	 * calling {@link SessionFactory#getNamedEntityGraphs(Class)} instead.
+	 * calling either {@link SessionFactory#getNamedEntityGraphs(Class)}
+	 * or {@link org.hibernate.query.SelectionQuery#getNamedEntityGraph(String)} instead.
 	 *
 	 * @see #find(EntityGraph, Object, FindOption...)
 	 * @see org.hibernate.query.SelectionQuery#setEntityGraph(EntityGraph, org.hibernate.graph.GraphSemantic)

hibernate-core/src/main/java/org/hibernate/procedure/internal/ProcedureCallImpl.java

@@ -19,6 +19,7 @@
 import java.util.Set;
 import java.util.stream.Stream;
 
+import jakarta.persistence.EntityGraph;
 import org.hibernate.HibernateException;
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
@@ -839,6 +840,10 @@ public Query<R> applyGraph(@SuppressWarnings("rawtypes") RootGraph graph, GraphS
 		throw new IllegalStateException( "EntityGraph hints are not supported for ProcedureCall/StoredProcedureQuery" );
 	}
 
+	@Override
+	public EntityGraph<? super R> getNamedEntityGraph(String graphName) {
+		throw new IllegalStateException( "EntityGraph hints are not supported for ProcedureCall/StoredProcedureQuery" );
+	}
 
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 	// outputs

hibernate-core/src/main/java/org/hibernate/query/Query.java

@@ -922,7 +922,7 @@ default Query<R> setPage(Page page) {
 	Query<R> setHint(String hintName, Object value);
 
 	@Override
-	Query<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic);
+	Query<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic);
 
 	@Override
 	Query<R> enableFetchProfile(String profileName);

hibernate-core/src/main/java/org/hibernate/query/SelectionQuery.java

@@ -291,7 +291,20 @@ default Stream<R> stream() {
 	 *
 	 * @since 6.3
 	 */
-	SelectionQuery<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic);
+	SelectionQuery<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic);
+
+	/**
+	 * Obtain an immutable reference to a predefined
+	 * {@linkplain jakarta.persistence.NamedEntityGraph named entity graph}
+	 * or return {@code null} if there is no predefined graph with the given
+	 * name.
+	 *
+	 * @param graphName The name of the predefined named entity graph
+	 *
+	 * @since 7.0
+	 */
+	@Incubating
+	EntityGraph<? super R> getNamedEntityGraph(String graphName);
 
 	/**
 	 * Enable the {@link org.hibernate.annotations.FetchProfile fetch profile}

hibernate-core/src/main/java/org/hibernate/query/spi/AbstractQuery.java

@@ -129,7 +129,7 @@ public QueryImplementor<R> setHint(String hintName, Object value) {
 	}
 
 	@Override
-	public QueryImplementor<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic) {
+	public QueryImplementor<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic) {
 		super.setEntityGraph( graph, semantic );
 		return this;
 	}

hibernate-core/src/main/java/org/hibernate/query/spi/AbstractSelectionQuery.java

@@ -374,11 +374,23 @@ public SelectionQuery<R> setHint(String hintName, Object value) {
 	}
 
 	@Override
-	public SelectionQuery<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic) {
-		applyGraph( (RootGraphImplementor<R>) graph, semantic );
+	public SelectionQuery<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic) {
+		applyGraph( (RootGraphImplementor<? super R>) graph, semantic );
 		return this;
 	}
 
+	@SuppressWarnings("unchecked")
+	@Override
+	public EntityGraph<? super R> getNamedEntityGraph(String graphName) {
+		for ( EntityGraph<? super R> graph : getSession().getSessionFactory()
+				.findEntityGraphsByType( (Class<R>) getResultType() ) ) {
+			if ( graph.getName().equals( graphName ) ) {
+				return graph;
+			}
+		}
+		return null;
+	}
+
 	@Override
 	public SelectionQuery<R> enableFetchProfile(String profileName) {
 		if ( getSession().getFactory().containsFetchProfileDefinition( profileName ) ) {

hibernate-core/src/main/java/org/hibernate/query/sql/internal/NativeQueryImpl.java

@@ -18,6 +18,7 @@
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+import jakarta.persistence.EntityGraph;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.hibernate.CacheMode;
 import org.hibernate.FlushMode;
@@ -596,6 +597,11 @@ public Query<R> applyGraph(@SuppressWarnings("rawtypes") RootGraph graph, GraphS
 		throw new HibernateException( "A native SQL query cannot use EntityGraphs" );
 	}
 
+	@Override
+	public EntityGraph<? super R> getNamedEntityGraph(String graphName) {
+		throw new HibernateException( "A native SQL query cannot use EntityGraphs" );
+	}
+
 	@Override
 	protected void applyEntityGraphHint(String hintName, Object value) {
 		super.applyEntityGraphHint( hintName, value );

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/QuerySqmImpl.java

@@ -811,7 +811,7 @@ public SqmQueryImplementor<R> setHint(String hintName, Object value) {
 	}
 
 	@Override
-	public Query<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic) {
+	public Query<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic) {
 		super.setEntityGraph( graph, semantic );
 		return this;
 	}

hibernate-core/src/main/java/org/hibernate/query/sqm/spi/DelegatingSqmSelectionQueryImplementor.java

@@ -157,11 +157,16 @@ public Optional<R> uniqueResultOptional() {
 	}
 
 	@Override
-	public SqmSelectionQueryImplementor<R> setEntityGraph(EntityGraph<R> graph, GraphSemantic semantic) {
+	public SqmSelectionQueryImplementor<R> setEntityGraph(EntityGraph<? super R> graph, GraphSemantic semantic) {
 		getDelegate().setEntityGraph( graph, semantic );
 		return this;
 	}
 
+	@Override
+	public EntityGraph<? super R> getNamedEntityGraph(String graphName) {
+		return getDelegate().getNamedEntityGraph( graphName );
+	}
+
 	@Override
 	public SqmSelectionQueryImplementor<R> enableFetchProfile(String profileName) {
 		getDelegate().enableFetchProfile( profileName );

hibernate-core/src/test/java/org/hibernate/orm/test/annotations/fetchprofile/NewGraphTest.java

@@ -8,11 +8,14 @@
 import org.hibernate.annotations.NaturalId;
 import org.hibernate.graph.GraphSemantic;
 import org.hibernate.graph.RootGraph;
+import org.hibernate.query.SelectionQuery;
 import org.hibernate.testing.orm.junit.DomainModel;
 import org.hibernate.testing.orm.junit.SessionFactory;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
 import org.junit.jupiter.api.Test;
 
+import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 import static jakarta.persistence.FetchType.LAZY;
@@ -21,7 +24,8 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @SessionFactory
-@DomainModel(annotatedClasses = {NewGraphTest.class, NewGraphTest.E.class, NewGraphTest.F.class, NewGraphTest.G.class, NewGraphTest.H.class})
+@DomainModel(annotatedClasses = {NewGraphTest.class, NewGraphTest.E.class, NewGraphTest.F.class, NewGraphTest.G.class, NewGraphTest.H.class,
+				NewGraphTest.A.class, NewGraphTest.Aa.class, NewGraphTest.B.class})
 public class NewGraphTest {
 
 	@Test void testByIdEntityGraph(SessionFactoryScope scope) {
@@ -168,6 +172,62 @@ public class NewGraphTest {
 		assertTrue( isInitialized( ee.f ) );
 	}
 
+	@Test
+	void subTypeEntityGraph(SessionFactoryScope scope) {
+		scope.inTransaction( s -> {
+			A a = new A();
+			Aa aa = new Aa();
+			B b1 = new B();
+			B b2 = new B();
+			B b3 = new B();
+			b1.a = a;
+			a.bs = new HashSet<>();
+			a.bs.add( b1 );
+
+			b2.a = aa;
+			aa.bs = new HashSet<>();
+			aa.bs.add( b2 );
+
+			b3.a = aa;
+			aa.bss = new HashSet<>();
+			aa.bss.add( b3 );
+
+			s.persist( a );
+			s.persist( aa );
+			s.persist( b1 );
+			s.persist( b2 );
+			s.persist( b3 );
+		} );
+
+		A a = scope.fromSession( s ->
+				s.createSelectionQuery( "from A", A.class )
+						.setMaxResults( 1 )
+						.getSingleResult() );
+		assertFalse( isInitialized( a.bs ) );
+
+		List<Aa> as = scope.fromSession( s -> {
+			SelectionQuery<Aa> query = s.createSelectionQuery( "from Aa", Aa.class );
+			return query
+					.setEntityGraph( query.getNamedEntityGraph( "a-graph" ), GraphSemantic.FETCH )
+					.getResultList();
+		} );
+		for ( Aa el : as ) {
+			assertTrue( isInitialized( el.bs ) );
+			assertFalse( isInitialized( el.bss ) );
+		}
+
+		as = scope.fromSession( s -> {
+			SelectionQuery<Aa> query = s.createSelectionQuery( "from Aa", Aa.class );
+			return query
+					.setEntityGraph( query.getNamedEntityGraph( "aa-graph" ), GraphSemantic.FETCH )
+					.getResultList();
+		} );
+		for ( Aa el : as ) {
+			assertTrue( isInitialized( el.bs ) );
+			assertTrue( isInitialized( el.bss ) );
+		}
+	}
+
 	@Entity(name = "E")
 	static class E {
 		@Id @GeneratedValue
@@ -199,4 +259,36 @@ static class H {
 		Long id;
 		@ManyToOne G g;
 	}
+
+	@NamedEntityGraph(name = "a-graph", attributeNodes = {@NamedAttributeNode("bs"),})
+	@Entity(name = "A")
+	static class A {
+		@Id @GeneratedValue
+		Long id;
+
+		@OneToMany(mappedBy = "a")
+		Set<B> bs;
+	}
+
+	@NamedEntityGraph(name = "aa-graph", attributeNodes = {@NamedAttributeNode("bs"), @NamedAttributeNode("bss"),})
+	@Entity(name = "Aa")
+	static class Aa extends A {
+		@Id @GeneratedValue
+		Long id;
+
+		@OneToMany(mappedBy = "aa")
+		Set<B> bss;
+	}
+
+	@Entity(name = "B")
+	static class B {
+		@Id @GeneratedValue
+		Long id;
+
+		@ManyToOne(fetch = LAZY)
+		A a;
+
+		@ManyToOne(fetch = LAZY)
+		Aa aa;
+	}
 }