HHH-18714 add support for inheritance in NamedEntityGraph annotation

Author: Asanio06

hibernate-core/src/main/java/org/hibernate/metamodel/model/domain/internal/JpaMetamodelImpl.java

@@ -20,6 +20,7 @@
 import java.util.stream.Collectors;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
+
 import org.hibernate.boot.model.NamedEntityGraphDefinition;
 import org.hibernate.boot.query.NamedQueryDefinition;
 import org.hibernate.boot.registry.classloading.spi.ClassLoaderService;
@@ -73,7 +74,6 @@
 import static org.hibernate.metamodel.internal.InjectionHelper.injectTypedQueryReference;
 
 /**
- *
  * @author Steve Ebersole
  */
 public class JpaMetamodelImpl implements JpaMetamodelImplementor, Serializable {
@@ -106,7 +106,7 @@ private ImportInfo(String importedName, Class<T> loadedClass) {
 	private final Map<Class<?>, String> entityProxyInterfaceMap = new HashMap<>();
 
 	private final Map<String, ImportInfo<?>> nameToImportMap = new ConcurrentHashMap<>();
-	private final Map<String,Object> knownInvalidnameToImportMap = new ConcurrentHashMap<>();
+	private final Map<String, Object> knownInvalidnameToImportMap = new ConcurrentHashMap<>();
 
 
 	public JpaMetamodelImpl(
@@ -143,13 +143,14 @@ public JpaCompliance getJpaCompliance() {
 	public <X> ManagedDomainType<X> managedType(String typeName) {
 		final ManagedDomainType<X> managedType = findManagedType( typeName );
 		if ( managedType == null ) {
-			throw new IllegalArgumentException("Not a managed type: " + typeName);
+			throw new IllegalArgumentException( "Not a managed type: " + typeName );
 		}
 		return managedType;
 	}
 
 	@Override
-	@Nullable public EntityDomainType<?> findEntityType(@Nullable String entityName) {
+	@Nullable
+	public EntityDomainType<?> findEntityType(@Nullable String entityName) {
 		if ( entityName == null ) {
 			return null;
 		}
@@ -162,18 +163,19 @@ public EntityDomainType<?> entity(String entityName) {
 		final EntityDomainType<?> entityType = findEntityType( entityName );
 		if ( entityType == null ) {
 			// per JPA
-			throw new IllegalArgumentException("Not an entity: " + entityName);
+			throw new IllegalArgumentException( "Not an entity: " + entityName );
 		}
 		return entityType;
 	}
 
 	@Override
-	@Nullable public EmbeddableDomainType<?> findEmbeddableType(@Nullable String embeddableName) {
+	@Nullable
+	public EmbeddableDomainType<?> findEmbeddableType(@Nullable String embeddableName) {
 		if ( embeddableName == null ) {
 			return null;
 		}
 		final ManagedDomainType<?> managedType = managedTypeByName.get( embeddableName );
-		if ( !( managedType instanceof EmbeddableDomainType<?> embeddableDomainType) ) {
+		if ( !( managedType instanceof EmbeddableDomainType<?> embeddableDomainType ) ) {
 			return null;
 		}
 		return embeddableDomainType;
@@ -183,7 +185,7 @@ public EntityDomainType<?> entity(String entityName) {
 	public EmbeddableDomainType<?> embeddable(String embeddableName) {
 		final EmbeddableDomainType<?> embeddableType = findEmbeddableType( embeddableName );
 		if ( embeddableType == null ) {
-			throw new IllegalArgumentException("Not an embeddable: " + embeddableName);
+			throw new IllegalArgumentException( "Not an embeddable: " + embeddableName );
 		}
 		return embeddableType;
 	}
@@ -226,7 +228,8 @@ public <X> EntityDomainType<X> resolveHqlEntityReference(String entityName) {
 	}
 
 	@Override
-	@Nullable public <X> ManagedDomainType<X> findManagedType(Class<X> cls) {
+	@Nullable
+	public <X> ManagedDomainType<X> findManagedType(Class<X> cls) {
 		//noinspection unchecked
 		return (ManagedDomainType<X>) managedTypeByClass.get( cls );
 	}
@@ -242,7 +245,8 @@ public <X> ManagedDomainType<X> managedType(Class<X> cls) {
 	}
 
 	@Override
-	@Nullable public <X> EntityDomainType<X> findEntityType(Class<X> cls) {
+	@Nullable
+	public <X> EntityDomainType<X> findEntityType(Class<X> cls) {
 		final ManagedType<?> type = managedTypeByClass.get( cls );
 		if ( !( type instanceof EntityDomainType<?> ) ) {
 			return null;
@@ -281,7 +285,7 @@ public <X> EmbeddableDomainType<X> embeddable(Class<X> cls) {
 
 	private Collection<ManagedDomainType<?>> getAllManagedTypes() {
 		// should never happen
-		return switch (jpaMetaModelPopulationSetting) {
+		return switch ( jpaMetaModelPopulationSetting ) {
 			case IGNORE_UNSUPPORTED -> managedTypeByClass.values();
 			case ENABLED -> managedTypeByName.values();
 			case DISABLED -> emptySet();
@@ -311,7 +315,7 @@ public Set<EmbeddableType<?>> getEmbeddables() {
 
 	@Override
 	public @Nullable Set<String> getEnumTypesForValue(String enumValue) {
-		return allowedEnumLiteralsToEnumTypeNames.get( enumValue);
+		return allowedEnumLiteralsToEnumTypeNames.get( enumValue );
 	}
 
 	@Override
@@ -388,7 +392,8 @@ public <T> void addNamedEntityGraph(String graphName, RootGraphImplementor<T> en
 		}
 	}
 
-	@Override @SuppressWarnings("unchecked")
+	@Override
+	@SuppressWarnings("unchecked")
 	public <T> RootGraphImplementor<T> findEntityGraphByName(String name) {
 		return (RootGraphImplementor<T>) entityGraphMap.get( name );
 	}
@@ -494,17 +499,48 @@ private void applyNamedEntityGraphs(Collection<NamedEntityGraphDefinition> named
 
 				);
 			}
+
 			final NamedEntityGraph namedEntityGraph = definition.getAnnotation();
 			final RootGraphImpl<?> entityGraph =
-					createRootGraph( definition.getRegisteredName(), entityType,
-							namedEntityGraph.includeAllAttributes() );
-			if ( namedEntityGraph.attributeNodes() != null ) {
-				applyNamedAttributeNodes( namedEntityGraph.attributeNodes(), namedEntityGraph, entityGraph );
-			}
+					createEntityGraph(
+							namedEntityGraph,
+							definition.getRegisteredName(),
+							entityType,
+							namedEntityGraph.includeAllAttributes()
+					);
+
 			entityGraphMap.put( definition.getRegisteredName(), entityGraph );
 		}
 	}
 
+	private <T> RootGraphImpl<T> createEntityGraph(
+			NamedEntityGraph namedEntityGraph,
+			String registeredName,
+			EntityDomainType<T> entityType,
+			boolean includeAllAttributes) {
+		final RootGraphImpl<T> entityGraph =
+				createRootGraph( registeredName, entityType, includeAllAttributes );
+
+		if ( namedEntityGraph.subclassSubgraphs() != null ) {
+			for ( var subclassSubgraph : namedEntityGraph.subclassSubgraphs() ) {
+				GraphImplementor<?> subgraph = (GraphImplementor<?>) entityGraph.addTreatedSubgraph(
+						(Class) subclassSubgraph.type() );
+
+				applyNamedAttributeNodes(
+						subclassSubgraph.attributeNodes(),
+						namedEntityGraph,
+						subgraph
+				);
+			}
+		}
+
+		if ( namedEntityGraph.attributeNodes() != null ) {
+			applyNamedAttributeNodes( namedEntityGraph.attributeNodes(), namedEntityGraph, entityGraph );
+		}
+
+		return entityGraph;
+	}
+
 	private static <T> RootGraphImpl<T> createRootGraph(
 			String name, EntityDomainType<T> entityType, boolean includeAllAttributes) {
 		final RootGraphImpl<T> entityGraph = new RootGraphImpl<>( name, entityType );
@@ -521,31 +557,44 @@ private void applyNamedAttributeNodes(
 			NamedEntityGraph namedEntityGraph,
 			GraphImplementor<?> graphNode) {
 		for ( NamedAttributeNode namedAttributeNode : namedAttributeNodes ) {
-			final AttributeNodeImplementor<?> attributeNode =
-					graphNode.findOrCreateAttributeNode( namedAttributeNode.value() );
+			final String value = namedAttributeNode.value();
+			final AttributeNodeImplementor<?> attributeNode = (AttributeNodeImplementor<?>) graphNode.addAttributeNode(
+					value );
+
 			if ( isNotEmpty( namedAttributeNode.subgraph() ) ) {
 				applyNamedSubgraphs(
 						namedEntityGraph,
 						namedAttributeNode.subgraph(),
-						attributeNode.makeSubGraph()
+						attributeNode,
+						false
 				);
 			}
 			if ( isNotEmpty( namedAttributeNode.keySubgraph() ) ) {
 				applyNamedSubgraphs(
 						namedEntityGraph,
 						namedAttributeNode.keySubgraph(),
-						attributeNode.makeKeySubGraph()
+						attributeNode,
+						true
 				);
 			}
 		}
 	}
 
-	private void applyNamedSubgraphs(
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	private <T> void applyNamedSubgraphs(
 			NamedEntityGraph namedEntityGraph,
 			String subgraphName,
-			SubGraphImplementor<?> subgraph) {
+			AttributeNodeImplementor<T> attributeNode, Boolean isKeySubGraph) {
 		for ( NamedSubgraph namedSubgraph : namedEntityGraph.subgraphs() ) {
 			if ( subgraphName.equals( namedSubgraph.name() ) ) {
+				final var isDefaultSubgraphType = namedSubgraph.type().equals( void.class );
+				final Class subGraphType = isDefaultSubgraphType ? null : namedSubgraph.type();
+
+				final SubGraphImplementor<?> subgraph = makeAttributeNodeSubgraph(
+						attributeNode, isKeySubGraph,
+						subGraphType
+				);
+
 				applyNamedAttributeNodes(
 						namedSubgraph.attributeNodes(),
 						namedEntityGraph,
@@ -555,6 +604,21 @@ private void applyNamedSubgraphs(
 		}
 	}
 
+	private static <T> SubGraphImplementor<?> makeAttributeNodeSubgraph(
+			AttributeNodeImplementor<T> attributeNode,
+			Boolean isKeySubGraph,
+			Class<T> subGraphType) {
+
+		if ( isKeySubGraph ) {
+			return subGraphType != null ? attributeNode.makeKeySubGraph( subGraphType )
+					: attributeNode.makeKeySubGraph();
+		}
+
+		return subGraphType != null ?
+				attributeNode.makeSubGraph( subGraphType ) :
+				attributeNode.makeSubGraph();
+	}
+
 	private <X> Class<X> resolveRequestedClass(String entityName) {
 		try {
 			return getServiceRegistry().requireService( ClassLoaderService.class )
@@ -631,7 +695,10 @@ public <T> EntityDomainType<T> resolveEntityReference(Class<T> javaType) {
 			}
 		}
 
-		throw new EntityTypeException( "Could not resolve entity class '" + javaType.getName() + "'", javaType.getName() );
+		throw new EntityTypeException(
+				"Could not resolve entity class '" + javaType.getName() + "'",
+				javaType.getName()
+		);
 	}
 
 	@Override
@@ -716,11 +783,30 @@ public void processJpa(
 
 	private void populateStaticMetamodel(MetadataImplementor bootMetamodel, MetadataContext context) {
 		bootMetamodel.visitNamedHqlQueryDefinitions( definition
-				-> injectTypedQueryReference( definition, namedQueryMetamodelClass( definition, context ) ) );
+															-> injectTypedQueryReference(
+				definition,
+				namedQueryMetamodelClass(
+						definition,
+						context
+				)
+		) );
 		bootMetamodel.visitNamedNativeQueryDefinitions( definition
-				-> injectTypedQueryReference( definition, namedQueryMetamodelClass( definition, context ) ) );
-		bootMetamodel.getNamedEntityGraphs().values().forEach(definition
-				-> injectEntityGraph( definition, graphMetamodelClass( definition, context ), this ) );
+																-> injectTypedQueryReference(
+				definition,
+				namedQueryMetamodelClass(
+						definition,
+						context
+				)
+		) );
+		bootMetamodel.getNamedEntityGraphs().values().forEach( definition
+																	-> injectEntityGraph(
+				definition,
+				graphMetamodelClass(
+						definition,
+						context
+				),
+				this
+		) );
 	}
 
 	private Class<?> namedQueryMetamodelClass(NamedQueryDefinition<?> definition, MetadataContext context) {
@@ -764,9 +850,7 @@ private <T> EntityDomainType<T> locateOrBuildEntityType(
 			MetadataContext context,
 			final TypeConfiguration typeConfiguration) {
 		@SuppressWarnings("unchecked")
-		final EntityDomainType<T> entityType =
-				(EntityDomainType<T>)
-						context.locateEntityType( persistentClass );
+		final EntityDomainType<T> entityType = (EntityDomainType<T>) context.locateEntityType( persistentClass );
 		return entityType == null
 				? buildEntityType( persistentClass, context, typeConfiguration )
 				: entityType;
@@ -818,8 +902,7 @@ private <T> MappedSuperclassDomainType<T> locateOrBuildMappedSuperclassType(
 			TypeConfiguration typeConfiguration) {
 		@SuppressWarnings("unchecked")
 		final MappedSuperclassDomainType<T> mappedSuperclassType =
-				(MappedSuperclassDomainType<T>)
-						context.locateMappedSuperclassType( mappedSuperclass );
+				(MappedSuperclassDomainType<T>) context.locateMappedSuperclassType( mappedSuperclass );
 		return mappedSuperclassType == null
 				? buildMappedSuperclassType( mappedSuperclass, context, typeConfiguration )
 				: mappedSuperclassType;