HHH-18178 Validation of CTE source query in `with` method

hibernate-core/src/main/java/org/hibernate/query/hql/internal/SemanticQueryBuilder.java

@@ -784,10 +784,7 @@ public class SemanticQueryBuilder<R> extends HqlParserBaseVisitor<Object> implem
 		final HqlParser.QueryExpressionContext queryExpressionContext = (HqlParser.QueryExpressionContext) ctx.getChild( queryExpressionIndex );
 		final SqmSelectQuery<Object> cte;
 		if ( cteContainer instanceof SqmSubQuery<?> ) {
-			cte = new SqmSubQuery<>(
+			cte = new SqmSubQuery<>( ( (SqmSubQuery<?>) cteContainer ).getParent(), creationContext.getNodeBuilder() );
-					processingStateStack.getCurrent().getProcessingQuery(),
-					creationContext.getNodeBuilder()
-			);
 		}
 		else {
 			cte = new SqmSelectStatement<>( creationContext.getNodeBuilder() );

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/AbstractSqmSelectQuery.java

@@ -168,7 +168,7 @@ public abstract class AbstractSqmSelectQuery<T>
 		return name;
 	}
 
-	private <X> JpaCteCriteria<X> withInternal(String name, AbstractQuery<X> criteria) {
+	protected <X> JpaCteCriteria<X> withInternal(String name, AbstractQuery<X> criteria) {
 		final SqmCteStatement<X> cteStatement = new SqmCteStatement<>(
 				name,
 				(SqmSelectQuery<X>) criteria,
@@ -181,7 +181,7 @@ public abstract class AbstractSqmSelectQuery<T>
 		return cteStatement;
 	}
 
-	private <X> JpaCteCriteria<X> withInternal(
+	protected <X> JpaCteCriteria<X> withInternal(
 			String name,
 			AbstractQuery<X> baseCriteria,
 			boolean unionDistinct,

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java

@@ -11,10 +11,12 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import org.hibernate.Internal;
 import org.hibernate.query.criteria.JpaCriteriaQuery;
+import org.hibernate.query.criteria.JpaCteCriteria;
 import org.hibernate.query.criteria.JpaExpression;
 import org.hibernate.query.criteria.JpaSelection;
 import org.hibernate.query.sqm.FetchClauseType;
@@ -32,6 +34,7 @@ import org.hibernate.query.sqm.tree.predicate.SqmPredicate;
 import org.hibernate.query.sqm.tree.from.SqmRoot;
 
 import jakarta.persistence.Tuple;
+import jakarta.persistence.criteria.AbstractQuery;
 import jakarta.persistence.criteria.Expression;
 import jakarta.persistence.criteria.Order;
 import jakarta.persistence.criteria.ParameterExpression;
@@ -252,6 +255,31 @@ public class SqmSelectStatement<T> extends AbstractSqmSelectQuery<T> implements
 		parameters.add( parameter );
 	}
 
+	@Override
+	protected <X> JpaCteCriteria<X> withInternal(String name, AbstractQuery<X> criteria) {
+		if ( criteria instanceof SqmSubQuery<?> ) {
+			throw new IllegalArgumentException(
+					"Invalid query type provided to root query 'with' method, " +
+							"expecting a root query to use as CTE instead found a subquery"
+			);
+		}
+		return super.withInternal( name, criteria );
+	}
+
+	@Override
+	protected <X> JpaCteCriteria<X> withInternal(
+			String name,
+			AbstractQuery<X> baseCriteria,
+			boolean unionDistinct,
+			Function<JpaCteCriteria<X>, AbstractQuery<X>> recursiveCriteriaProducer) {
+		if ( baseCriteria instanceof SqmSubQuery<?> ) {
+			throw new IllegalArgumentException(
+					"Invalid query type provided to root query 'with' method, " +
+							"expecting a root query to use as CTE instead found a subquery"
+			);
+		}
+		return super.withInternal( name, baseCriteria, unionDistinct, recursiveCriteriaProducer );
+	}
 
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 	// JPA

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSubQuery.java

@@ -14,6 +14,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import org.hibernate.metamodel.model.domain.EntityDomainType;
 import org.hibernate.query.criteria.JpaCrossJoin;
@@ -201,6 +202,32 @@ public class SqmSubQuery<T> extends AbstractSqmSelectQuery<T> implements SqmSele
 		return ( (JpaCteContainer) parent ).getCteCriteria( cteName );
 	}
 
+	@Override
+	protected <X> JpaCteCriteria<X> withInternal(String name, AbstractQuery<X> criteria) {
+		if ( !( criteria instanceof SqmSubQuery<?> ) || ( (SqmSubQuery<X>) criteria ).getParent() != parent ) {
+			throw new IllegalArgumentException(
+					"Invalid query type provided to subquery 'with' method, " +
+							"expecting a subquery with the same parent to use as CTE"
+			);
+		}
+		return super.withInternal( name, criteria );
+	}
+
+	@Override
+	protected <X> JpaCteCriteria<X> withInternal(
+			String name,
+			AbstractQuery<X> baseCriteria,
+			boolean unionDistinct,
+			Function<JpaCteCriteria<X>, AbstractQuery<X>> recursiveCriteriaProducer) {
+		if ( !( baseCriteria instanceof SqmSubQuery<?> ) || ( (SqmSubQuery<X>) baseCriteria ).getParent() != parent ) {
+			throw new IllegalArgumentException(
+					"Invalid query type provided to subquery 'with' method, " +
+							"expecting a subquery with the same parent to use as CTE"
+			);
+		}
+		return super.withInternal( name, baseCriteria, unionDistinct, recursiveCriteriaProducer );
+	}
+
 	@Override
 	public SqmQuery<?> getContainingQuery() {
 		return parent;