HHH-18816 FK rendering: check parent query specs for group by / order by

Author: mbellade

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

@@ -21,6 +21,7 @@
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.internal.util.collections.CollectionHelper;
+import org.hibernate.internal.util.collections.Stack;
 import org.hibernate.jpa.spi.JpaCompliance;
 import org.hibernate.metamodel.mapping.BasicValuedMapping;
 import org.hibernate.metamodel.mapping.Bindable;
@@ -168,48 +169,72 @@ public static ModelPartContainer getTargetMappingIfNeeded(
 			SqmPath<?> sqmPath,
 			ModelPartContainer modelPartContainer,
 			SqmToSqlAstConverter sqlAstCreationState) {
-		final SqmQueryPart<?> queryPart = sqlAstCreationState.getCurrentSqmQueryPart();
-		if ( queryPart != null ) {
-			// We only need to do this for queries
-			final Clause clause = sqlAstCreationState.getCurrentClauseStack().getCurrent();
-			if ( clause != Clause.FROM && modelPartContainer.getPartMappingType() != modelPartContainer && sqmPath.getLhs() instanceof SqmFrom<?, ?> ) {
-				final ModelPart modelPart;
-				if ( modelPartContainer instanceof PluralAttributeMapping ) {
-					modelPart = getCollectionPart(
-							(PluralAttributeMapping) modelPartContainer,
-							castNonNull( sqmPath.getNavigablePath().getParent() )
-					);
-				}
-				else {
-					modelPart = modelPartContainer;
-				}
-				if ( modelPart instanceof EntityAssociationMapping association ) {
-					// If the path is one of the association's target key properties,
-					// we need to render the target side if in group/order by
-					if ( association.getTargetKeyPropertyNames().contains( sqmPath.getReferencedPathSource().getPathName() )
-							&& ( clause == Clause.GROUP || clause == Clause.ORDER
-							|| !isFkOptimizationAllowed( sqmPath.getLhs(), association )
-							|| queryPart.getFirstQuerySpec().groupByClauseContains( sqmPath.getNavigablePath(), sqlAstCreationState )
-							|| queryPart.getFirstQuerySpec().orderByClauseContains( sqmPath.getNavigablePath(), sqlAstCreationState ) ) ) {
-						return association.getAssociatedEntityMappingType();
-					}
+		// We only need to do this for queries
+		if ( sqlAstCreationState.getCurrentClauseStack().getCurrent() != Clause.FROM
+				&& modelPartContainer.getPartMappingType() != modelPartContainer && sqmPath.getLhs() instanceof SqmFrom<?, ?> ) {
+			final ModelPart modelPart;
+			if ( modelPartContainer instanceof PluralAttributeMapping pluralAttributeMapping ) {
+				modelPart = getCollectionPart(
+						pluralAttributeMapping,
+						castNonNull( sqmPath.getNavigablePath().getParent() )
+				);
+			}
+			else {
+				modelPart = modelPartContainer;
+			}
+			if ( modelPart instanceof EntityAssociationMapping association ) {
+				if ( shouldRenderTargetSide( sqmPath, association, sqlAstCreationState ) ) {
+					return association.getAssociatedEntityMappingType();
 				}
 			}
 		}
 		return modelPartContainer;
 	}
 
+	private static boolean shouldRenderTargetSide(
+			SqmPath<?> sqmPath,
+			EntityAssociationMapping association,
+			SqmToSqlAstConverter sqlAstCreationState) {
+		if ( !association.getTargetKeyPropertyNames().contains( sqmPath.getReferencedPathSource().getPathName() ) ) {
+			return false;
+		}
+		// If the path is one of the association's target key properties,
+		// we need to render the target side if in group/order by
+		final Clause clause = sqlAstCreationState.getCurrentClauseStack().getCurrent();
+		return clause == Clause.GROUP || clause == Clause.ORDER
+				|| !isFkOptimizationAllowed( sqmPath.getLhs(), association )
+				|| clauseContainsPath( Clause.GROUP, sqmPath, sqlAstCreationState )
+				|| clauseContainsPath( Clause.ORDER, sqmPath, sqlAstCreationState );
+	}
+
+	private static boolean clauseContainsPath(
+			Clause clauseToCheck,
+			SqmPath<?> sqmPath,
+			SqmToSqlAstConverter sqlAstCreationState) {
+		final Stack<SqmQueryPart> queryPartStack = sqlAstCreationState.getSqmQueryPartStack();
+		final NavigablePath navigablePath = sqmPath.getNavigablePath();
+		final Boolean found = queryPartStack.findCurrentFirst( queryPart -> {
+			final SqmQuerySpec<?> spec = queryPart.getFirstQuerySpec();
+			if ( clauseToCheck == Clause.GROUP && spec.groupByClauseContains( navigablePath, sqlAstCreationState ) ) {
+				return true;
+			}
+			else if ( clauseToCheck == Clause.ORDER && spec.orderByClauseContains( navigablePath, sqlAstCreationState ) ) {
+				return true;
+			}
+			else {
+				return null;
+			}
+		} );
+		return Boolean.TRUE.equals( found );
+	}
+
 	private static CollectionPart getCollectionPart(PluralAttributeMapping attribute, NavigablePath path) {
 		final CollectionPart.Nature nature = CollectionPart.Nature.fromNameExact( path.getLocalName() );
-		if ( nature != null ) {
-			switch ( nature ) {
-				case ELEMENT:
-					return attribute.getElementDescriptor();
-				case INDEX:
-					return attribute.getIndexDescriptor();
-			}
-		}
-		return null;
+		return nature != null ? switch ( nature ) {
+			case ELEMENT -> attribute.getElementDescriptor();
+			case INDEX -> attribute.getIndexDescriptor();
+			case ID -> attribute.getIdentifierDescriptor();
+		} : null;
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java

@@ -488,7 +488,7 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 	private boolean deduplicateSelectionItems;
 	private ForeignKeyDescriptor.Nature currentlyResolvingForeignKeySide;
 	private SqmStatement<?> currentSqmStatement;
-	private SqmQueryPart<?> currentSqmQueryPart;
+	private Stack<SqmQueryPart> sqmQueryPartStack = new StandardStack<>( SqmQueryPart.class );
 	private CteContainer cteContainer;
 	/**
 	 * A map from {@link SqmCteTable#getCteName()} to the final SQL name.
@@ -792,9 +792,10 @@ public Stack<Clause> getCurrentClauseStack() {
 		return currentClauseStack;
 	}
 
+	@SuppressWarnings("rawtypes")
 	@Override
-	public SqmQueryPart<?> getCurrentSqmQueryPart() {
-		return currentSqmQueryPart;
+	public Stack<SqmQueryPart> getSqmQueryPartStack() {
+		return sqmQueryPartStack;
 	}
 
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1726,8 +1727,7 @@ public CteStatement visitCteStatement(SqmCteStatement<?> sqmCteStatement) {
 							);
 							final DelegatingSqmAliasedNodeCollector collector =
 									(DelegatingSqmAliasedNodeCollector) processingState.getSqlExpressionResolver();
-							final SqmQueryPart<?> oldSqmQueryPart = currentSqmQueryPart;
-							currentSqmQueryPart = queryGroup;
+							sqmQueryPartStack.push( queryGroup );
 							pushProcessingState( processingState );
 
 							try {
@@ -1771,7 +1771,7 @@ public CteStatement visitCteStatement(SqmCteStatement<?> sqmCteStatement) {
 							}
 							finally {
 								popProcessingStateStack();
-								currentSqmQueryPart = oldSqmQueryPart;
+								sqmQueryPartStack.pop();
 							}
 						}
 						finally {
@@ -1985,8 +1985,7 @@ public QueryGroup visitQueryGroup(SqmQueryGroup<?> queryGroup) {
 		);
 		final DelegatingSqmAliasedNodeCollector collector = (DelegatingSqmAliasedNodeCollector) processingState
 				.getSqlExpressionResolver();
-		final SqmQueryPart<?> sqmQueryPart = currentSqmQueryPart;
-		currentSqmQueryPart = queryGroup;
+		sqmQueryPartStack.push( queryGroup );
 		pushProcessingState( processingState );
 
 		try {
@@ -2008,7 +2007,7 @@ public QueryGroup visitQueryGroup(SqmQueryGroup<?> queryGroup) {
 		}
 		finally {
 			popProcessingStateStack();
-			currentSqmQueryPart = sqmQueryPart;
+			sqmQueryPartStack.pop();
 		}
 	}
 
@@ -2043,9 +2042,8 @@ public QuerySpec visitQuerySpec(SqmQuerySpec<?> sqmQuerySpec) {
 			);
 		}
 
-		final SqmQueryPart<?> sqmQueryPart = currentSqmQueryPart;
 		final boolean originalDeduplicateSelectionItems = deduplicateSelectionItems;
-		currentSqmQueryPart = sqmQuerySpec;
+		sqmQueryPartStack.push( sqmQuerySpec );
 		// In sub-queries, we can never deduplicate the selection items as that might change semantics
 		deduplicateSelectionItems = false;
 		pushProcessingState( processingState );
@@ -2062,7 +2060,7 @@ public QuerySpec visitQuerySpec(SqmQuerySpec<?> sqmQuerySpec) {
 			inNestedContext = oldInNestedContext;
 			popProcessingStateStack();
 			queryTransformers.pop();
-			currentSqmQueryPart = sqmQueryPart;
+			sqmQueryPartStack.pop();
 			deduplicateSelectionItems = originalDeduplicateSelectionItems;
 		}
 	}
@@ -2203,7 +2201,7 @@ public SelectClause visitSelectClause(SqmSelectClause selectClause) {
 		try {
 			final SelectClause sqlSelectClause = currentQuerySpec().getSelectClause();
 			if ( selectClause == null ) {
-				final SqmFrom<?, ?> implicitSelection = determineImplicitSelection( (SqmQuerySpec<?>) currentSqmQueryPart );
+				final SqmFrom<?, ?> implicitSelection = determineImplicitSelection( (SqmQuerySpec<?>) getCurrentSqmQueryPart() );
 				visitSelection( 0, new SqmSelection<>( implicitSelection, implicitSelection.nodeBuilder() ) );
 			}
 			else {
@@ -2228,7 +2226,7 @@ public SelectClause visitSelectClause(SqmSelectClause selectClause) {
 	@Override
 	public Void visitSelection(SqmSelection<?> sqmSelection) {
 		visitSelection(
-				currentSqmQueryPart.getFirstQuerySpec().getSelectClause().getSelections().indexOf( sqmSelection ),
+				getCurrentSqmQueryPart().getFirstQuerySpec().getSelectClause().getSelections().indexOf( sqmSelection ),
 				sqmSelection
 		);
 		return null;
@@ -2353,7 +2351,7 @@ else if ( statement.getQuerySource() == SqmQuerySource.CRITERIA && currentClause
 			// To avoid this issue, we determine the position and let the SqlAstTranslator handle the rest.
 			// Usually it will render `select ?, count(*) from dual group by 1` if supported
 			// or force rendering the parameter as literal instead so that the database can see the grouping is fine
-			final SqmQuerySpec<?> querySpec = currentSqmQueryPart.getFirstQuerySpec();
+			final SqmQuerySpec<?> querySpec = getCurrentSqmQueryPart().getFirstQuerySpec();
 			sqmPosition = indexOfExpression( querySpec.getSelectClause().getSelections(), groupByClauseExpression );
 			path = null;
 		}
@@ -2395,7 +2393,7 @@ private SqlSelectionExpression selectionExpression(List<SqlSelection> selections
 				return null;
 			}
 		}
-		return currentSqmQueryPart instanceof SqmQueryGroup<?>
+		return getCurrentSqmQueryPart() instanceof SqmQueryGroup<?>
 				// Reusing the SqlSelection for query groups would be wrong because the aliases do no exist
 				// So we have to use a literal expression in a new SqlSelection instance to refer to the position
 				? sqlSelectionExpression( selection )

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/FakeSqmToSqlAstConverter.java

@@ -92,6 +92,11 @@ public Stack<Clause> getCurrentClauseStack() {
 		return null;
 	}
 
+	@Override
+	public Stack<SqmQueryPart> getSqmQueryPartStack() {
+		return null;
+	}
+
 	@Override
 	public SqmQueryPart<?> getCurrentSqmQueryPart() {
 		return null;

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/SqmToSqlAstConverter.java

@@ -33,7 +33,11 @@
 public interface SqmToSqlAstConverter extends SemanticQueryWalker<Object>, SqlAstCreationState {
 	Stack<Clause> getCurrentClauseStack();
 
-	SqmQueryPart<?> getCurrentSqmQueryPart();
+	Stack<SqmQueryPart> getSqmQueryPartStack();
+
+	default SqmQueryPart<?> getCurrentSqmQueryPart() {
+		return getSqmQueryPartStack().getCurrent();
+	}
 
 	void registerQueryTransformer(QueryTransformer transformer);