CountQuery creation and getProjection(...) with JSqlParser.

Inside `JSqlParserQueryUtils` it is now possible to create count queries with `createCountQueryFor(...)`.
Furthermore, the `getProjection(...)` method utilizes the power of the JSQlParser.

Related tickets spring-projects#2409

Author: DiegoKrupitza

src/main/java/org/springframework/data/jpa/repository/query/JSqlParserQueryUtils.java

@@ -1,9 +1,10 @@
 package org.springframework.data.jpa.repository.query;
 
+import static org.springframework.data.jpa.repository.query.QueryUtils.checkSortExpression;
+
 import java.util.*;
 import java.util.stream.Collectors;
 
-import net.sf.jsqlparser.statement.select.*;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Order;
 import org.springframework.data.util.Streamable;
@@ -22,12 +23,12 @@
 import net.sf.jsqlparser.parser.CCJSqlParserUtil;
 import net.sf.jsqlparser.schema.Column;
 import net.sf.jsqlparser.schema.Table;
+import net.sf.jsqlparser.statement.select.*;
 import net.sf.jsqlparser.util.SelectUtils;
 
-import static org.springframework.data.jpa.repository.query.QueryUtils.checkSortExpression;
-
 /**
- * Simple utility class to create JPA queries using JSqlParser.
+ * Simple utility class to create SQL queries using JSqlParser. The usage of this class is limited to valid SQL syntax.
+ * If the query or the statement is not valid any usage of this class will throw an exception.
  *
  * @author Diego Krupitza
  */
@@ -95,7 +96,17 @@ public static String getQueryString(String template, String entityName) {
 	public static String detectAlias(String query) {
 		Select selectStatement = parseSelectStatement(query);
 		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+		return detectAlias(selectBody);
+	}
 
+	/**
+	 * Resolves the alias for the entity to be retrieved from the given {@link PlainSelect}. Note that you only provide
+	 * valid Query strings. Things such as <code>from User u</code> will throw an {@link IllegalArgumentException}.
+	 *
+	 * @param selectBody must not be {@literal null}.
+	 * @return Might return {@literal null}.
+	 */
+	private static String detectAlias(PlainSelect selectBody) {
 		Alias alias = selectBody.getFromItem().getAlias();
 		return alias == null ? null : alias.getName();
 	}
@@ -262,6 +273,96 @@ private static OrderByElement getOrderClause(final Set<String> joinAliases, fina
 		return orderByElement;
 	}
 
+	/**
+	 * Creates a count projected query from the given original query.
+	 *
+	 * @param originalQuery must not be {@literal null} or empty.
+	 * @return Guaranteed to be not {@literal null}.
+	 */
+	public static String createCountQueryFor(String originalQuery) {
+		return createCountQueryFor(originalQuery, null);
+	}
+
+	/**
+	 * Creates a count projected query from the given original query.
+	 *
+	 * @param originalQuery must not be {@literal null}.
+	 * @param countProjection may be {@literal null}.
+	 * @return a query String to be used a count query for pagination. Guaranteed to be not {@literal null}.
+	 */
+	public static String createCountQueryFor(String originalQuery, @Nullable String countProjection) {
+
+		Assert.hasText(originalQuery, "OriginalQuery must not be null or empty!");
+
+		Select selectStatement = parseSelectStatement(originalQuery);
+		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+
+		// remove order by
+		selectBody.setOrderByElements(null);
+
+		if (StringUtils.hasText(countProjection)) {
+			Function jSqlCount = getJSqlCount(Collections.singletonList(countProjection), false);
+			selectBody.setSelectItems(Collections.singletonList(new SelectExpressionItem(jSqlCount)));
+			return selectBody.toString();
+		}
+
+		boolean distinct = selectBody.getDistinct() != null;
+		selectBody.setDistinct(null); // reset possible distinct
+
+		String tableAlias = detectAlias(selectBody);
+
+		// is never null
+		List<SelectItem> selectItems = selectBody.getSelectItems();
+
+		if (onlyASingleColumnProjection(selectItems)) {
+			SelectExpressionItem singleProjection = (SelectExpressionItem) selectItems.get(0);
+
+			Column column = (Column) singleProjection.getExpression();
+			String countProp = column.getFullyQualifiedName();
+
+			Function jSqlCount = getJSqlCount(Collections.singletonList(countProp), distinct);
+			selectBody.setSelectItems(Collections.singletonList(new SelectExpressionItem(jSqlCount)));
+			return selectBody.toString();
+		}
+
+		String countProp = tableAlias == null ? "*" : tableAlias;
+
+		Function jSqlCount = getJSqlCount(Collections.singletonList(countProp), distinct);
+		selectBody.setSelectItems(Collections.singletonList(new SelectExpressionItem(jSqlCount)));
+
+		return selectBody.toString();
+	}
+
+	/**
+	 * Checks whether a given projection only contains a single column definition (aka without functions, etc)
+	 * 
+	 * @param projection the projection to analyse
+	 * @return <code>true</code> when the projection only contains a single column definition otherwise <code>false</code>
+	 */
+	private static boolean onlyASingleColumnProjection(List<SelectItem> projection) {
+		// this is unfortunately the only way to check without any hacky & hard string regex magic
+		return projection.size() == 1 && projection.get(0) instanceof SelectExpressionItem
+				&& (((SelectExpressionItem) projection.get(0)).getExpression()) instanceof Column;
+	}
+
+	/**
+	 * Returns the projection part of the query, i.e. everything between {@code select} and {@code from}.
+	 *
+	 * @param query must not be {@literal null} or empty.
+	 * @return the projection part of the query.
+	 */
+	public static String getProjection(String query) {
+		Assert.hasText(query, "Query must not be null or empty!");
+
+		Select selectStatement = parseSelectStatement(query);
+		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+
+		return selectBody.getSelectItems() //
+				.stream() //
+				.map(Object::toString) //
+				.collect(Collectors.joining(", ")).trim();
+	}
+
 	/**
 	 * Generates a JSqlParser table from an entity name and an optional alias name
 	 * 

src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryUtilsUnitTests.java

@@ -1,11 +1,12 @@
 package org.springframework.data.jpa.repository.query;
 
 import static org.assertj.core.api.Assertions.*;
-import static org.springframework.data.jpa.repository.query.QueryUtils.getOuterJoinAliases;
+import static org.springframework.data.jpa.repository.query.QueryUtils.detectAlias;
 
 import java.util.Collections;
 import java.util.Set;
 
+import org.assertj.core.api.SoftAssertions;
 import org.junit.jupiter.api.Test;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.domain.Sort;
@@ -58,11 +59,6 @@ void detectAliasThrowsErrorOnInvalidSQL() {
 				.isInstanceOf(IllegalArgumentException.class);
 	}
 
-	@Test
-	void allowsFullyQualifiedEntityNamesInQuery() {
-		assertThat(JSqlParserQueryUtils.detectAlias(FQ_QUERY)).isEqualTo("u");
-	}
-
 	@Test // DATAJPA-798
 	void detectsAliasInQueryContainingLineBreaks() {
 		assertThat(JSqlParserQueryUtils.detectAlias("select \n u \n from \n User \nu")).isEqualTo("u");
@@ -347,6 +343,161 @@ void discoversCorrectAliasForJoinFetch() {
 		assertThat(aliases).containsExactly("authority");
 	}
 
+	@Test // DATAJPA-420
+	void createsCountQueryForScalarSelects() {
+		assertCountQuery("select p.lastname,p.firstname from Person p", "select count(p) from Person p");
+	}
+
+	@Test // DATAJPA-456
+	void createCountQueryFromTheGivenCountProjection() {
+		assertThat(JSqlParserQueryUtils.createCountQueryFor("select p.lastname,p.firstname from Person p", "p.lastname"))
+				.isEqualToIgnoringCase("select count(p.lastname) from Person p");
+	}
+
+	@Test // DATAJPA-736
+	void supportsNonAsciiCharactersInEntityNames() {
+		assertThat(JSqlParserQueryUtils.createCountQueryFor("select u from Usèr u"))
+				.isEqualToIgnoringCase("select count(u) from Usèr u");
+	}
+
+	@Test // DATAJPA-1500
+	void createCountQuerySupportsWhitespaceCharacters() {
+
+		assertThat(JSqlParserQueryUtils.createCountQueryFor("select * from User user\n" + //
+				"  where user.age = 18\n" + //
+				"  order by user.name\n ")).isEqualToIgnoringCase("select count(user) from User user where user.age = 18");
+	}
+
+	@Test
+	void createCountQuerySupportsLineBreaksInSelectClause() {
+
+		assertThat(JSqlParserQueryUtils.createCountQueryFor("select user.age,\n" + //
+				"  user.name\n" + //
+				"  from User user\n" + //
+				"  where user.age = 18\n" + //
+				"  order\nby\nuser.name\n ")).isEqualToIgnoringCase("select count(user) from User user where user.age = 18");
+	}
+
+	@Test
+	void createCountQuerySupportsLineBreakRightAfterDistinct() {
+
+		assertThat(JSqlParserQueryUtils.createCountQueryFor("select\ndistinct\nuser.age,\n" + //
+				"user.name\n" + //
+				"from\nUser\nuser")).isEqualTo(JSqlParserQueryUtils.createCountQueryFor("select\ndistinct user.age,\n" + //
+						"user.name\n" + //
+						"from\nUser\nuser"));
+	}
+
+	@Test
+	void createsCountQueryCorrectly() {
+		assertCountQuery(QUERY, COUNT_QUERY);
+	}
+
+	@Test
+	void createsCountQueriesCorrectlyForCapitalLetter() {
+		assertCountQuery("SELECT u FROM User u where u.foo.bar = ?", "select count(u) FROM User u where u.foo.bar = ?");
+	}
+
+	@Test
+	void createsCountQueryForDistinctQueries() {
+
+		assertCountQuery("select distinct u from User u where u.foo = ?",
+				"select count(distinct u) from User u where u.foo = ?");
+	}
+
+	@Test
+	void failsOnConstructorQueries() {
+
+		assertThatThrownBy(
+				() -> JSqlParserQueryUtils.createCountQueryFor("select distinct new User(u.name) from User u where u.foo = ?"))
+						.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	void createsCountQueryForJoins() {
+
+		assertCountQuery("select distinct u.name from User u left outer join u.roles r WHERE r = ?",
+				"select count(distinct u.name) from User u left outer join u.roles r WHERE r = ?");
+	}
+
+	@Test
+	void createsCountQueryForQueriesWithSubSelects() {
+
+		assertCountQuery("select u from User u left outer join u.roles r where r in (select r from Role)",
+				"select count(u) from User u left outer join u.roles r where r in (select r from Role)");
+	}
+
+	@Test
+	void createsCountQueryForAliasesCorrectly() {
+
+		assertCountQuery("select u from User as u", "select count(u) from User as u");
+	}
+
+	@Test
+	void doesNotAllowShortJpaSyntax() {
+		assertThatThrownBy(() -> assertCountQuery(INVALID_QUERY, COUNT_QUERY)).isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	void allowsFullyQualifiedEntityNamesInQuery() {
+
+		assertThat(detectAlias(FQ_QUERY)).isEqualTo("u");
+		assertCountQuery(FQ_QUERY, "select count(u) from org.acme.domain.User$Foo_Bar u");
+	}
+
+	@Test // DATAJPA-342
+	void usesReturnedVariableInCOuntProjectionIfSet() {
+
+		assertCountQuery("select distinct m.genre from Media m where m.user = ?1 order by m.genre asc",
+				"select count(distinct m.genre) from Media m where m.user = ?1");
+	}
+
+	@Test // DATAJPA-343
+	void projectsCOuntQueriesForQueriesWithSubselects() {
+
+		assertCountQuery("select o from Foo o where cb.id in (select b from Bar b)",
+				"select count(o) from Foo o where cb.id in (select b from Bar b)");
+	}
+
+	@Test // DATAJPA-377
+	void removesOrderByInGeneratedCountQueryFromOriginalQueryIfPresent() {
+
+		assertCountQuery("select distinct m.genre from Media m where m.user = ?1 OrDer  By   m.genre ASC",
+				"select count(distinct m.genre) from Media m where m.user = ?1");
+	}
+
+	@Test // DATAJPA-409
+	void createsCountQueryForNestedReferenceCorrectly() {
+		assertCountQuery("select a.b from A a", "select count(a.b) from A a");
+	}
+
+	@Test // DATAJPA-1679
+	void findProjectionClauseWithDistinct() {
+
+		SoftAssertions.assertSoftly(sofly -> {
+			sofly.assertThat(JSqlParserQueryUtils.getProjection("select * from x")).isEqualTo("*");
+			sofly.assertThat(JSqlParserQueryUtils.getProjection("select a, b, c from x")).isEqualTo("a, b, c");
+			sofly.assertThat(JSqlParserQueryUtils.getProjection("select distinct a, b, c from x")).isEqualTo("a, b, c");
+			sofly.assertThat(JSqlParserQueryUtils.getProjection("select DISTINCT a, b, c from x")).isEqualTo("a, b, c");
+		});
+	}
+
+	@Test // DATAJPA-1696
+	void findProjectionClauseWithSubselect() {
+
+		// This is not a required behavior the testcase in QueryUtilsUnitTests#findProjectionClauseWithSubselect tells why
+		assertThat(JSqlParserQueryUtils.getProjection("select * from (select x from y)")).isEqualTo("*");
+	}
+
+	@Test // DATAJPA-1696
+	void findProjectionClauseWithIncludedFrom() {
+		assertThat(JSqlParserQueryUtils.getProjection("select x, frommage, y from t")).isEqualTo("x, frommage, y");
+	}
+
+	private static void assertCountQuery(String originalQuery, String countQuery) {
+		assertThat(JSqlParserQueryUtils.createCountQueryFor(originalQuery)).isEqualToIgnoringCase(countQuery);
+	}
+
 	private static void endsIgnoringCase(String original, String endWithIgnoreCase) {
 		// https://github.com/assertj/assertj-core/pull/2451
 		// can be removed when upgrading to version 3.23.0 assertJ