Hacking.

Author: mp911de

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaCountQueryCreator.java

@@ -15,16 +15,12 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import jakarta.persistence.criteria.CriteriaBuilder;
-import jakarta.persistence.criteria.CriteriaQuery;
-import jakarta.persistence.criteria.Expression;
-import jakarta.persistence.criteria.Predicate;
-import jakarta.persistence.criteria.Root;
+import jakarta.persistence.EntityManager;
 
 import org.springframework.data.domain.Sort;
+import org.springframework.data.jpa.repository.support.JpqlQueryTemplates;
 import org.springframework.data.repository.query.ReturnedType;
 import org.springframework.data.repository.query.parser.PartTree;
-import org.springframework.lang.Nullable;
 
 /**
  * Special {@link JpaQueryCreator} that creates a count projecting query.
@@ -36,41 +32,35 @@
  */
 public class JpaCountQueryCreator extends JpaQueryCreator {
 
-	private boolean distinct;
+	private final boolean distinct;
+	private final ReturnedType returnedType;
 
 	/**
-	 * Creates a new {@link JpaCountQueryCreator}.
+	 * Creates a new {@link JpaCountQueryCreator}
 	 *
 	 * @param tree
-	 * @param type
-	 * @param builder
+	 * @param returnedType
 	 * @param provider
+	 * @param templates
+	 * @param em
 	 */
-	public JpaCountQueryCreator(PartTree tree, ReturnedType type, CriteriaBuilder builder,
-			ParameterMetadataProvider provider) {
+	public JpaCountQueryCreator(PartTree tree, ReturnedType returnedType, ParameterMetadataProvider provider,
+			JpqlQueryTemplates templates, EntityManager em) {
 
-		super(tree, type, builder, provider);
+		super(tree, returnedType, provider, templates, em);
 
 		this.distinct = tree.isDistinct();
+		this.returnedType = returnedType;
 	}
 
 	@Override
-	protected CriteriaQuery<? extends Object> createCriteriaQuery(CriteriaBuilder builder, ReturnedType type) {
+	protected JpqlQueryBuilder.AbstractJpqlQuery buildQuery(Sort sort) {
 
-		return builder.createQuery(Long.class);
-	}
-
-	@Override
-	@SuppressWarnings("unchecked")
-	protected CriteriaQuery<? extends Object> complete(@Nullable Predicate predicate, Sort sort,
-			CriteriaQuery<? extends Object> query, CriteriaBuilder builder, Root<?> root) {
-
-		CriteriaQuery<? extends Object> select = query.select(getCountQuery(query, builder, root));
-		return predicate == null ? select : select.where(predicate);
-	}
+		JpqlQueryBuilder.SelectStep selectStep = JpqlQueryBuilder.selectFrom(returnedType.getDomainType());
+		if (this.distinct) {
+			selectStep = selectStep.distinct();
+		}
 
-	@SuppressWarnings("rawtypes")
-	private Expression getCountQuery(CriteriaQuery<?> query, CriteriaBuilder builder, Root<?> root) {
-		return distinct ? builder.countDistinct(root) : builder.count(root);
+		return selectStep.count();
 	}
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaKeysetScrollQueryCreator.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.jpa.repository.query;
 
+import jakarta.persistence.EntityManager;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.Predicate;
@@ -27,6 +28,7 @@
 import org.springframework.data.domain.KeysetScrollPosition;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.repository.support.JpaEntityInformation;
+import org.springframework.data.jpa.repository.support.JpqlQueryTemplates;
 import org.springframework.data.repository.query.ReturnedType;
 import org.springframework.data.repository.query.parser.PartTree;
 import org.springframework.lang.Nullable;
@@ -42,21 +44,21 @@ class JpaKeysetScrollQueryCreator extends JpaQueryCreator {
 	private final JpaEntityInformation<?, ?> entityInformation;
 	private final KeysetScrollPosition scrollPosition;
 
-	public JpaKeysetScrollQueryCreator(PartTree tree, ReturnedType type, CriteriaBuilder builder,
-			ParameterMetadataProvider provider, JpaEntityInformation<?, ?> entityInformation,
-			KeysetScrollPosition scrollPosition) {
+	public JpaKeysetScrollQueryCreator(PartTree tree, ReturnedType type, ParameterMetadataProvider provider,
+			JpqlQueryTemplates templates, JpaEntityInformation<?, ?> entityInformation, KeysetScrollPosition scrollPosition,
+			EntityManager em) {
 
-		super(tree, type, builder, provider);
+		super(tree, type, provider, templates, em);
 
 		this.entityInformation = entityInformation;
 		this.scrollPosition = scrollPosition;
 	}
 
-	@Override
+	// TODO
 	protected CriteriaQuery<?> complete(@Nullable Predicate predicate, Sort sort, CriteriaQuery<?> query,
 			CriteriaBuilder builder, Root<?> root) {
 
-		KeysetScrollSpecification<Object> keysetSpec = new KeysetScrollSpecification<>(scrollPosition, sort,
+		/*KeysetScrollSpecification<Object> keysetSpec = new KeysetScrollSpecification<>(scrollPosition, sort,
 				entityInformation);
 		Predicate keysetPredicate = keysetSpec.createPredicate(root, builder);
 
@@ -69,7 +71,10 @@ protected CriteriaQuery<?> complete(@Nullable Predicate predicate, Sort sort, Cr
 			return queryToUse.where(keysetPredicate);
 		}
 
-		return queryToUse;
+		return queryToUse;*/
+
+		// TODO
+		return null;
 	}
 
 	@Override