Introduce QueryTransformer for functions

Author: beikov

hibernate-core/src/main/java/org/hibernate/query/sqm/function/SelfRenderingSqmAggregateFunction.java

@@ -21,6 +21,7 @@
 import org.hibernate.query.sqm.tree.predicate.SqmPredicate;
 import org.hibernate.query.sqm.tree.select.SqmSelectableNode;
 import org.hibernate.sql.ast.tree.SqlAstNode;
+import org.hibernate.sql.ast.tree.expression.Expression;
 import org.hibernate.sql.ast.tree.predicate.Predicate;
 
 /**
@@ -74,7 +75,7 @@ public SelfRenderingSqmAggregateFunction<T> copy(SqmCopyContext context) {
 	}
 
 	@Override
-	public SelfRenderingFunctionSqlAstExpression convertToSqlAst(SqmToSqlAstConverter walker) {
+	public Expression convertToSqlAst(SqmToSqlAstConverter walker) {
 		final ReturnableType<?> resultType = resolveResultType(
 				walker.getCreationContext().getMappingMetamodel().getTypeConfiguration()
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/function/SelfRenderingSqmFunction.java

@@ -23,6 +23,7 @@
 import org.hibernate.query.sqm.tree.SqmVisitableNode;
 import org.hibernate.query.sqm.tree.expression.SqmFunction;
 import org.hibernate.sql.ast.tree.SqlAstNode;
+import org.hibernate.sql.ast.tree.expression.Expression;
 import org.hibernate.type.spi.TypeConfiguration;
 
 import static java.util.Collections.emptyList;
@@ -109,7 +110,7 @@ protected static List<SqlAstNode> resolveSqlAstArguments(List<? extends SqmTyped
 	}
 
 	@Override
-	public SelfRenderingFunctionSqlAstExpression convertToSqlAst(SqmToSqlAstConverter walker) {
+	public Expression convertToSqlAst(SqmToSqlAstConverter walker) {
 		final ReturnableType<?> resultType = resolveResultType(
 				walker.getCreationContext().getMappingMetamodel().getTypeConfiguration()
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/function/SelfRenderingSqmOrderedSetAggregateFunction.java

@@ -25,6 +25,7 @@
 import org.hibernate.query.sqm.tree.select.SqmSortSpecification;
 import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.tree.SqlAstNode;
+import org.hibernate.sql.ast.tree.expression.Expression;
 import org.hibernate.sql.ast.tree.predicate.Predicate;
 import org.hibernate.sql.ast.tree.select.SortSpecification;
 
@@ -91,7 +92,7 @@ public SelfRenderingSqmOrderedSetAggregateFunction<T> copy(SqmCopyContext contex
 	}
 
 	@Override
-	public SelfRenderingFunctionSqlAstExpression convertToSqlAst(SqmToSqlAstConverter walker) {
+	public Expression convertToSqlAst(SqmToSqlAstConverter walker) {
 		final ReturnableType<?> resultType = resolveResultType(
 				walker.getCreationContext().getMappingMetamodel().getTypeConfiguration()
 		);

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

@@ -94,28 +94,28 @@
 import org.hibernate.persister.entity.AbstractEntityPersister;
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.persister.entity.Joinable;
-import org.hibernate.query.sqm.BinaryArithmeticOperator;
 import org.hibernate.query.BindableType;
-import org.hibernate.query.sqm.CastType;
-import org.hibernate.query.sqm.ComparisonOperator;
-import org.hibernate.query.sqm.DynamicInstantiationNature;
-import org.hibernate.query.sqm.FetchClauseType;
-import org.hibernate.query.spi.NavigablePath;
 import org.hibernate.query.QueryLogging;
 import org.hibernate.query.ReturnableType;
 import org.hibernate.query.SemanticException;
-import org.hibernate.query.sqm.SortOrder;
-import org.hibernate.query.sqm.TemporalUnit;
-import org.hibernate.query.sqm.UnaryArithmeticOperator;
 import org.hibernate.query.criteria.JpaPath;
+import org.hibernate.query.spi.NavigablePath;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.query.spi.QueryParameterBinding;
 import org.hibernate.query.spi.QueryParameterBindings;
 import org.hibernate.query.spi.QueryParameterImplementor;
+import org.hibernate.query.sqm.BinaryArithmeticOperator;
+import org.hibernate.query.sqm.CastType;
+import org.hibernate.query.sqm.ComparisonOperator;
+import org.hibernate.query.sqm.DynamicInstantiationNature;
+import org.hibernate.query.sqm.FetchClauseType;
 import org.hibernate.query.sqm.InterpretationException;
+import org.hibernate.query.sqm.SortOrder;
 import org.hibernate.query.sqm.SqmExpressible;
 import org.hibernate.query.sqm.SqmPathSource;
 import org.hibernate.query.sqm.SqmQuerySource;
+import org.hibernate.query.sqm.TemporalUnit;
+import org.hibernate.query.sqm.UnaryArithmeticOperator;
 import org.hibernate.query.sqm.function.AbstractSqmSelfRenderingFunctionDescriptor;
 import org.hibernate.query.sqm.function.SelfRenderingAggregateFunctionSqlAstExpression;
 import org.hibernate.query.sqm.function.SelfRenderingFunctionSqlAstExpression;
@@ -260,6 +260,7 @@
 import org.hibernate.sql.ast.spi.SqlSelection;
 import org.hibernate.sql.ast.tree.Statement;
 import org.hibernate.sql.ast.tree.cte.CteColumn;
+import org.hibernate.sql.ast.tree.cte.CteContainer;
 import org.hibernate.sql.ast.tree.cte.CteStatement;
 import org.hibernate.sql.ast.tree.cte.CteTable;
 import org.hibernate.sql.ast.tree.cte.SearchClauseSpecification;
@@ -286,6 +287,7 @@
 import org.hibernate.sql.ast.tree.expression.Over;
 import org.hibernate.sql.ast.tree.expression.Overflow;
 import org.hibernate.sql.ast.tree.expression.QueryLiteral;
+import org.hibernate.sql.ast.tree.expression.QueryTransformer;
 import org.hibernate.sql.ast.tree.expression.SelfRenderingExpression;
 import org.hibernate.sql.ast.tree.expression.SelfRenderingSqlFragmentExpression;
 import org.hibernate.sql.ast.tree.expression.SqlSelectionExpression;
@@ -383,6 +385,7 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 	private final SqlAstCreationContext creationContext;
 	private final boolean jpaQueryComplianceEnabled;
 	private final SqmStatement<?> statement;
+	private final CteContainer cteContainer = new GlobalCteContainer();
 
 	private final QueryOptions queryOptions;
 	private final LoadQueryInfluencers loadQueryInfluencers;
@@ -416,6 +419,7 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 	private final Stack<Supplier<MappingModelExpressible<?>>> inferrableTypeAccessStack = new StandardStack<>(
 			() -> null
 	);
+	private final Stack<List<QueryTransformer>> queryTransformers = new StandardStack<>();
 	private boolean inTypeInference;
 
 	private SqmByUnit appliedByUnit;
@@ -645,7 +649,7 @@ public Statement visitStatement(SqmStatement<?> sqmStatement) {
 
 	@Override
 	public UpdateStatement visitUpdateStatement(SqmUpdateStatement<?> sqmStatement) {
-		Map<String, CteStatement> cteStatements = this.visitCteContainer( sqmStatement );
+		final CteContainer cteContainer = this.visitCteContainer( sqmStatement );
 
 		final SqmRoot<?> sqmTarget = sqmStatement.getTarget();
 		final String entityName = sqmTarget.getEntityName();
@@ -711,7 +715,7 @@ public UpdateStatement visitUpdateStatement(SqmUpdateStatement<?> sqmStatement)
 			}
 
 			return new UpdateStatement(
-					sqmStatement.isWithRecursive(), cteStatements,
+					cteContainer,
 					(NamedTableReference) rootTableGroup.getPrimaryTableReference(),
 					assignments,
 					SqlAstTreeHelper.combinePredicates( suppliedPredicate, additionalRestrictions ),
@@ -889,7 +893,7 @@ public Expression resolveSqlExpression(
 
 	@Override
 	public DeleteStatement visitDeleteStatement(SqmDeleteStatement<?> statement) {
-		Map<String, CteStatement> cteStatements = this.visitCteContainer( statement );
+		final CteContainer cteContainer = this.visitCteContainer( statement );
 
 		final String entityName = statement.getTarget().getEntityName();
 		final EntityPersister entityDescriptor = creationContext.getSessionFactory()
@@ -945,8 +949,7 @@ public DeleteStatement visitDeleteStatement(SqmDeleteStatement<?> statement) {
 			}
 
 			return new DeleteStatement(
-					statement.isWithRecursive(),
-					cteStatements,
+					cteContainer,
 					(NamedTableReference) rootTableGroup.getPrimaryTableReference(),
 					SqlAstTreeHelper.combinePredicates( suppliedPredicate, additionalRestrictions ),
 					Collections.emptyList()
@@ -962,7 +965,7 @@ public DeleteStatement visitDeleteStatement(SqmDeleteStatement<?> statement) {
 
 	@Override
 	public InsertStatement visitInsertSelectStatement(SqmInsertSelectStatement<?> sqmStatement) {
-		Map<String, CteStatement> cteStatements = this.visitCteContainer( sqmStatement );
+		final CteContainer cteContainer = this.visitCteContainer( sqmStatement );
 
 		final String entityName = sqmStatement.getTarget().getEntityName();
 		final EntityPersister entityDescriptor = creationContext.getSessionFactory()
@@ -1003,8 +1006,7 @@ public InsertStatement visitInsertSelectStatement(SqmInsertSelectStatement<?> sq
 			getFromClauseAccess().registerTableGroup( rootPath, rootTableGroup );
 
 			insertStatement = new InsertStatement(
-					sqmStatement.isWithRecursive(),
-					cteStatements,
+					cteContainer,
 					(NamedTableReference) rootTableGroup.getPrimaryTableReference(),
 					Collections.emptyList()
 			);
@@ -1049,7 +1051,7 @@ public InsertStatement visitInsertSelectStatement(SqmInsertSelectStatement<?> sq
 
 	@Override
 	public InsertStatement visitInsertValuesStatement(SqmInsertValuesStatement<?> sqmStatement) {
-		Map<String, CteStatement> cteStatements = this.visitCteContainer( sqmStatement );
+		final CteContainer cteContainer = this.visitCteContainer( sqmStatement );
 		final String entityName = sqmStatement.getTarget().getEntityName();
 		final EntityPersister entityDescriptor = creationContext.getSessionFactory()
 				.getRuntimeMetamodels()
@@ -1085,8 +1087,7 @@ public InsertStatement visitInsertValuesStatement(SqmInsertValuesStatement<?> sq
 			getFromClauseAccess().registerTableGroup( rootPath, rootTableGroup );
 
 			final InsertStatement insertStatement = new InsertStatement(
-					sqmStatement.isWithRecursive(),
-					cteStatements,
+					cteContainer,
 					(NamedTableReference) rootTableGroup.getPrimaryTableReference(),
 					Collections.emptyList()
 			);
@@ -1369,10 +1370,10 @@ public Values visitValues(SqmValues sqmValues) {
 
 	@Override
 	public SelectStatement visitSelectStatement(SqmSelectStatement<?> statement) {
-		Map<String, CteStatement> cteStatements = this.visitCteContainer( statement );
+		final CteContainer cteContainer = this.visitCteContainer( statement );
 		final QueryPart queryPart = visitQueryPart( statement.getQueryPart() );
 		final List<DomainResult<?>> domainResults = queryPart.isRoot() ? this.domainResults : Collections.emptyList();
-		return new SelectStatement( statement.isWithRecursive(), cteStatements, queryPart, domainResults );
+		return new SelectStatement( cteContainer, queryPart, domainResults );
 	}
 
 	@Override
@@ -1558,14 +1559,15 @@ public static CteTable createCteTable(
 	}
 
 	@Override
-	public Map<String, CteStatement> visitCteContainer(SqmCteContainer consumer) {
+	public CteContainer visitCteContainer(SqmCteContainer consumer) {
 		final Collection<SqmCteStatement<?>> sqmCteStatements = consumer.getCteStatements();
-		final Map<String, CteStatement> cteStatements = new LinkedHashMap<>( sqmCteStatements.size() );
+		if ( consumer.isWithRecursive() ) {
+			cteContainer.setWithRecursive( true );
+		}
 		for ( SqmCteStatement<?> sqmCteStatement : sqmCteStatements ) {
-			final CteStatement cteStatement = visitCteStatement( sqmCteStatement );
-			cteStatements.put( cteStatement.getCteTable().getTableExpression(), cteStatement );
+			cteContainer.addCteStatement( visitCteStatement( sqmCteStatement ) );
 		}
-		return cteStatements;
+		return cteContainer;
 	}
 
 	private boolean trackSelectionsForGroup;
@@ -1635,7 +1637,7 @@ public QuerySpec visitQuerySpec(SqmQuerySpec<?> sqmQuerySpec) {
 		);
 		final SqmSelectClause selectClause = sqmQuerySpec.getSelectClause();
 
-		Predicate originalAdditionalRestrictions = additionalRestrictions;
+		final Predicate originalAdditionalRestrictions = additionalRestrictions;
 		additionalRestrictions = null;
 
 		final boolean trackAliasedNodePositions;
@@ -1686,6 +1688,7 @@ else if ( sqmQuerySpec.hasPositionalGroupItem() ) {
 		// In sub-queries, we can never deduplicate the selection items as that might change semantics
 		deduplicateSelectionItems = false;
 		pushProcessingState( processingState );
+		queryTransformers.push( new ArrayList<>() );
 
 		try {
 			// we want to visit the from-clause first
@@ -1725,14 +1728,23 @@ else if ( sqmQuerySpec.hasPositionalGroupItem() ) {
 				applyCollectionFilterPredicates( sqlQuerySpec );
 			}
 
-			return sqlQuerySpec;
+			QuerySpec finalQuerySpec = sqlQuerySpec;
+			for ( QueryTransformer transformer : queryTransformers.getCurrent() ) {
+				finalQuerySpec = transformer.transform(
+						cteContainer,
+						finalQuerySpec,
+						this
+				);
+			}
+			return finalQuerySpec;
 		}
 		finally {
 			if ( additionalRestrictions != null ) {
 				sqlQuerySpec.applyPredicate( additionalRestrictions );
 			}
 			additionalRestrictions = originalAdditionalRestrictions;
 			popProcessingStateStack();
+			queryTransformers.pop();
 			currentSqmQueryPart = sqmQueryPart;
 			deduplicateSelectionItems = originalDeduplicateSelectionItems;
 		}
@@ -4566,6 +4578,11 @@ public Expression visitFunction(SqmFunction<?> sqmFunction) {
 		}
 	}
 
+	@Override
+	public void registerQueryTransformer(QueryTransformer transformer) {
+		queryTransformers.getCurrent().add( transformer );
+	}
+
 	@Override
 	public Star visitStar(SqmStar sqmStar) {
 		return new Star();
@@ -6313,4 +6330,38 @@ private static JdbcMappingContainer highestPrecedence(JdbcMappingContainer type1
 
 		return type1;
 	}
+
+	private class GlobalCteContainer implements CteContainer {
+		private final Map<String, CteStatement> cteStatements;
+		private boolean recursive;
+
+		public GlobalCteContainer() {
+			this.cteStatements = new LinkedHashMap<>();
+		}
+
+		@Override
+		public boolean isWithRecursive() {
+			return recursive;
+		}
+
+		@Override
+		public void setWithRecursive(boolean recursive) {
+			this.recursive = recursive;
+		}
+
+		@Override
+		public Map<String, CteStatement> getCteStatements() {
+			return cteStatements;
+		}
+
+		@Override
+		public CteStatement getCteStatement(String cteLabel) {
+			return cteStatements.get( cteLabel );
+		}
+
+		@Override
+		public void addCteStatement(CteStatement cteStatement) {
+			cteStatements.put( cteStatement.getCteTable().getTableExpression(), cteStatement );
+		}
+	}
 }

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

@@ -20,6 +20,7 @@
 import org.hibernate.sql.ast.spi.SqlAstProcessingState;
 import org.hibernate.sql.ast.spi.SqlExpressionResolver;
 import org.hibernate.sql.ast.tree.expression.Expression;
+import org.hibernate.sql.ast.tree.expression.QueryTransformer;
 
 /**
  *
@@ -73,6 +74,10 @@ public Stack<Clause> getCurrentClauseStack() {
 		return null;
 	}
 
+	@Override
+	public void registerQueryTransformer(QueryTransformer transformer) {
+	}
+
 	@Override
 	public List<Expression> expandSelfRenderingFunctionMultiValueParameter(SqmParameter<?> sqmParameter) {
 		return null;

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

@@ -14,6 +14,7 @@
 import org.hibernate.sql.ast.spi.SqlAstCreationState;
 import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.tree.expression.Expression;
+import org.hibernate.sql.ast.tree.expression.QueryTransformer;
 
 /**
  * Specialized SemanticQueryWalker (SQM visitor) for producing SQL AST.
@@ -23,6 +24,8 @@
 public interface SqmToSqlAstConverter extends SemanticQueryWalker<Object>, SqlAstCreationState {
 	Stack<Clause> getCurrentClauseStack();
 
+	void registerQueryTransformer(QueryTransformer transformer);
+
 	List<Expression> expandSelfRenderingFunctionMultiValueParameter(SqmParameter<?> sqmParameter);
 
 }

hibernate-core/src/main/java/org/hibernate/sql/ast/spi/AbstractSqlAstTranslator.java

@@ -1313,7 +1313,7 @@ protected void visitReturningColumns(MutationStatement mutationStatement) {
 	}
 
 	public void visitCteContainer(CteContainer cteContainer) {
-		final Collection<CteStatement> cteStatements = cteContainer.getCteStatements();
+		final Collection<CteStatement> cteStatements = cteContainer.getCteStatements().values();
 		if ( cteStatements.isEmpty() ) {
 			return;
 		}