HHH-17410 Support creating count query from existing query

Author: beikov

hibernate-core/src/main/java/org/hibernate/query/criteria/CriteriaDefinition.java

@@ -400,4 +400,9 @@ public <X> JpaDerivedRoot<X> from(Subquery<X> subquery) {
 	public <X> JpaRoot<X> from(JpaCteCriteria<X> cte) {
 		return query.from(cte);
 	}
+
+	@Override
+	public JpaCriteriaQuery<Long> createCountQuery() {
+		return query.createCountQuery();
+	}
 }

hibernate-core/src/main/java/org/hibernate/query/criteria/JpaCriteriaQuery.java

@@ -26,6 +26,13 @@
  */
 public interface JpaCriteriaQuery<T> extends CriteriaQuery<T>, JpaQueryableCriteria<T>, JpaSelectCriteria<T> {
 
+	/**
+	 * Wraps this query in a subquery and returns a count query based on that subquery in the from clause.
+	 *
+	 * @since 6.4
+	 */
+	JpaCriteriaQuery<Long> createCountQuery();
+
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 	// Limit/Offset/Fetch clause
 

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

@@ -6,6 +6,7 @@
  */
 package org.hibernate.query.sqm.tree;
 
+import org.hibernate.Incubating;
 import org.hibernate.query.sqm.internal.NoParamSqmCopyContext;
 import org.hibernate.query.sqm.internal.SimpleSqmCopyContext;
 
@@ -18,6 +19,16 @@ public interface SqmCopyContext {
 
 	<T> T registerCopy(T original, T copy);
 
+	/**
+	 * Returns whether the {@code fetch} flag for attribute joins should be copied over.
+	 *
+	 * @since 6.4
+	 */
+	@Incubating
+	default boolean copyFetchedFlag() {
+		return true;
+	}
+
 	static SqmCopyContext simpleContext() {
 		return new SimpleSqmCopyContext();
 	}

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/domain/SqmBagJoin.java

@@ -65,7 +65,7 @@ public SqmBagJoin<O, E> copy(SqmCopyContext context) {
 						getAttribute(),
 						getExplicitAlias(),
 						getSqmJoinType(),
-						isFetched(),
+						context.copyFetchedFlag() && isFetched(),
 						nodeBuilder()
 				)
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/domain/SqmListJoin.java

@@ -68,7 +68,7 @@ public SqmListJoin<O, E> copy(SqmCopyContext context) {
 						getAttribute(),
 						getExplicitAlias(),
 						getSqmJoinType(),
-						isFetched(),
+						context.copyFetchedFlag() && isFetched(),
 						nodeBuilder()
 				)
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/domain/SqmMapJoin.java

@@ -8,7 +8,6 @@
 
 import java.util.Map;
 import jakarta.persistence.criteria.Expression;
-import jakarta.persistence.criteria.Path;
 import jakarta.persistence.criteria.Predicate;
 
 import org.hibernate.metamodel.model.domain.EntityDomainType;
@@ -67,7 +66,7 @@ public SqmMapJoin<O, K, V> copy(SqmCopyContext context) {
 						getAttribute(),
 						getExplicitAlias(),
 						getSqmJoinType(),
-						isFetched(),
+						context.copyFetchedFlag() && isFetched(),
 						nodeBuilder()
 				)
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/domain/SqmSetJoin.java

@@ -67,7 +67,7 @@ public SqmSetJoin<O, E> copy(SqmCopyContext context) {
 						getModel(),
 						getExplicitAlias(),
 						getSqmJoinType(),
-						isFetched(),
+						context.copyFetchedFlag() && isFetched(),
 						nodeBuilder()
 				)
 		);

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/domain/SqmSingularJoin.java

@@ -11,7 +11,6 @@
 import org.hibernate.metamodel.model.domain.EntityDomainType;
 import org.hibernate.metamodel.model.domain.SingularPersistentAttribute;
 import org.hibernate.query.sqm.SemanticQueryWalker;
-import org.hibernate.query.sqm.SqmPathSource;
 import org.hibernate.spi.NavigablePath;
 import org.hibernate.query.hql.spi.SqmCreationProcessingState;
 import org.hibernate.query.sqm.NodeBuilder;
@@ -76,7 +75,7 @@ public SqmSingularJoin<O, T> copy(SqmCopyContext context) {
 						getAttribute(),
 						getExplicitAlias(),
 						getSqmJoinType(),
-						isFetched(),
+						context.copyFetchedFlag() && isFetched(),
 						nodeBuilder()
 				)
 		);

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

@@ -6,6 +6,7 @@
  */
 package org.hibernate.query.sqm.tree.select;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -27,10 +28,12 @@
 import org.hibernate.query.sqm.NodeBuilder;
 import org.hibernate.query.sqm.SemanticQueryWalker;
 import org.hibernate.query.sqm.SqmQuerySource;
+import org.hibernate.query.sqm.internal.NoParamSqmCopyContext;
 import org.hibernate.query.sqm.internal.SqmUtil;
 import org.hibernate.query.sqm.tree.SqmCopyContext;
 import org.hibernate.query.sqm.tree.SqmStatement;
 import org.hibernate.query.sqm.tree.cte.SqmCteStatement;
+import org.hibernate.query.sqm.tree.expression.SqmStar;
 import org.hibernate.query.sqm.tree.expression.ValueBindJpaCriteriaParameter;
 import org.hibernate.query.sqm.tree.expression.SqmParameter;
 import org.hibernate.query.sqm.tree.from.SqmFromClause;
@@ -455,4 +458,83 @@ private void validateComplianceFetchOffset() {
 							"Please disable the JPA query compliance if you want to use this feature." );
 		}
 	}
+
+	@Override
+	public JpaCriteriaQuery<Long> createCountQuery() {
+		final SqmCopyContext context = new NoParamSqmCopyContext() {
+			@Override
+			public boolean copyFetchedFlag() {
+				return false;
+			}
+		};
+		final NodeBuilder nodeBuilder = nodeBuilder();
+		final Set<SqmParameter<?>> parameters;
+		if ( this.parameters == null ) {
+			parameters = null;
+		}
+		else {
+			parameters = new LinkedHashSet<>( this.parameters.size() );
+			for ( SqmParameter<?> parameter : this.parameters ) {
+				parameters.add( parameter.copy( context ) );
+			}
+		}
+		final SqmSelectStatement<Long> selectStatement = new SqmSelectStatement<>(
+				nodeBuilder,
+				copyCteStatements( context ),
+				Long.class,
+				SqmQuerySource.CRITERIA,
+				parameters
+		);
+		final SqmQuerySpec<Long> querySpec = new SqmQuerySpec<>( nodeBuilder );
+
+		final SqmSubQuery<Tuple> subquery = new SqmSubQuery<>( selectStatement, Tuple.class, nodeBuilder );
+		final SqmQueryPart<T> queryPart = getQueryPart().copy( context );
+		resetSelections( queryPart );
+		// Reset the
+		if ( queryPart.getFetch() == null && queryPart.getOffset() == null ) {
+			queryPart.setOrderByClause( null );
+		}
+		//noinspection unchecked
+		subquery.setQueryPart( (SqmQueryPart<Tuple>) queryPart );
+
+		querySpec.setFromClause( new SqmFromClause( 1 ) );
+		querySpec.setSelectClause( new SqmSelectClause( false, 1, nodeBuilder ) );
+		selectStatement.setQueryPart( querySpec );
+		selectStatement.select( nodeBuilder.count( new SqmStar( nodeBuilder ) ) );
+		selectStatement.from( subquery );
+		return selectStatement;
+	}
+
+	private void resetSelections(SqmQueryPart<?> queryPart) {
+		if ( queryPart instanceof SqmQuerySpec<?> ) {
+			resetSelections( (SqmQuerySpec<?>) queryPart );
+		}
+		else {
+			final SqmQueryGroup<?> group = (SqmQueryGroup<?>) queryPart;
+			for ( SqmQueryPart<?> part : group.getQueryParts() ) {
+				resetSelections( part );
+			}
+		}
+	}
+
+	private void resetSelections(SqmQuerySpec<?> querySpec) {
+		final NodeBuilder nodeBuilder = nodeBuilder();
+		final List<SqmSelection<?>> selections = querySpec.getSelectClause().getSelections();
+		final List<SqmSelectableNode<?>> subSelections = new ArrayList<>();
+
+		if ( selections.isEmpty() ) {
+			subSelections.add( (SqmSelectableNode<?>) nodeBuilder.literal( 1 ).alias( "c0" ) );
+		}
+		else {
+			for ( SqmSelection<?> selection : selections ) {
+				selection.getSelectableNode().visitSubSelectableNodes( e -> {
+					e.alias( "c" + subSelections.size() );
+					subSelections.add( e );
+				} );
+			}
+		}
+
+		querySpec.getSelectClause().setSelection( (SqmSelectableNode<?>) nodeBuilder.tuple( subSelections ) );
+	}
+
 }