Remove JSqlParserQueryUtils in favour of JSqlParserQueryEnhancer.

Since all the methods within `JSqlParserQueryUtils` were only used by `JSqlParserQueryEnhancer` it makes more sense to just implement it there. Additionally we introduced the `JSqlParserUtils` class which is there for operations that a commonly used when working with JSqlParser (for example generating a count function object).

Related tickets spring-projects#2409

Author: DiegoKrupitza

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

@@ -17,6 +17,8 @@
 
 import org.springframework.data.domain.Sort;
 
+import java.util.Set;
+
 /**
  * The implementation of {@link QueryEnhancer} using {@link QueryUtils}.
  *
@@ -60,6 +62,11 @@ public String getProjection() {
 		return QueryUtils.getProjection(this.query.getQueryString());
 	}
 
+	@Override
+	public Set<String> getJoinAliases() {
+		return QueryUtils.getOuterJoinAliases(this.query.getQueryString());
+	}
+
 	@Override
 	public DeclaredQuery getQuery() {
 		return this.query;

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

@@ -15,15 +15,38 @@
  */
 package org.springframework.data.jpa.repository.query;
 
+import net.sf.jsqlparser.JSQLParserException;
+import net.sf.jsqlparser.expression.Alias;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.expression.Function;
+import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
+import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
+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 org.springframework.data.domain.Sort;
+import org.springframework.data.util.Streamable;
+import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+import static org.springframework.data.jpa.repository.query.JSqlParserUtils.*;
+import static org.springframework.data.jpa.repository.query.QueryUtils.checkSortExpression;
 
 /**
- * The implementation of {@link QueryEnhancer} using {@link JSqlParserQueryUtils}.
+ * The implementation of {@link QueryEnhancer} using JSqlParser.
  *
  * @author Diego Krupitza
  */
 public class JSqlParserQueryEnhancer implements QueryEnhancer {
 
+	private static final String DEFAULT_TABLE_ALIAS = "x";
+
 	private final DeclaredQuery query;
 
 	/**
@@ -35,32 +58,275 @@ public JSqlParserQueryEnhancer(DeclaredQuery query) {
 
 	@Override
 	public String getExistsQueryString(String entityName, String countQueryPlaceHolder, Iterable<String> idAttributes) {
-		return JSqlParserQueryUtils.getExistsQueryString(entityName, countQueryPlaceHolder, idAttributes);
+		final Table tableNameWithAlias = getTableWithAlias(entityName, DEFAULT_TABLE_ALIAS);
+		Function jSqlCount = getJSqlCount(Collections.singletonList(countQueryPlaceHolder), false);
+
+		Select select = SelectUtils.buildSelectFromTableAndSelectItems(tableNameWithAlias,
+				new SelectExpressionItem(jSqlCount));
+
+		PlainSelect selectBody = (PlainSelect) select.getSelectBody();
+
+		List<Expression> equalityExpressions = Streamable.of(idAttributes).stream() //
+				.map(field -> {
+					Expression tableNameField = new Column().withTable(tableNameWithAlias).withColumnName(field);
+					Expression inputField = new Column(":".concat(field));
+					return new EqualsTo(tableNameField, inputField);
+				}).collect(Collectors.toList());
+
+		if (equalityExpressions.size() > 1) {
+			AndExpression rootOfWhereClause = concatenateWithAndExpression(equalityExpressions);
+			selectBody.setWhere(rootOfWhereClause);
+		} else if (equalityExpressions.size() == 1) {
+			selectBody.setWhere(equalityExpressions.get(0));
+		}
+
+		return selectBody.toString();
 	}
 
 	@Override
 	public String getQueryString(String template, String entityName) {
-		return JSqlParserQueryUtils.getQueryString(template, entityName);
+		Assert.hasText(entityName, "Entity name must not be null or empty!");
+		return String.format(template, entityName);
 	}
 
 	@Override
 	public String applySorting(Sort sort, String alias) {
-		return JSqlParserQueryUtils.applySorting(this.query.getQueryString(), sort, alias);
+		String queryString = query.getQueryString();
+		Assert.hasText(queryString, "Query must not be null or empty!");
+
+		if (sort.isUnsorted()) {
+			return queryString;
+		}
+
+		Select selectStatement = parseSelectStatement(queryString);
+		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+
+		final Set<String> joinAliases = getJoinAliases(selectBody);
+
+		final Set<String> selectionAliases = getSelectionAliases(selectBody);
+
+		List<OrderByElement> orderByElements = sort.stream() //
+				.map(order -> getOrderClause(joinAliases, selectionAliases, alias, order)) //
+				.collect(Collectors.toList());
+
+		if (CollectionUtils.isEmpty(selectBody.getOrderByElements())) {
+			selectBody.setOrderByElements(new ArrayList<>());
+		}
+
+		selectBody.getOrderByElements().addAll(orderByElements);
+
+		return selectBody.toString();
+
+	}
+
+	/**
+	 * Returns the aliases used inside the selection part in the query.
+	 *
+	 * @param selectBody a {@link PlainSelect} containing a query. Must not be {@literal null}.
+	 * @return a {@literal Set} containing all found aliases. Guaranteed to be not {@literal null}.
+	 */
+	private Set<String> getSelectionAliases(PlainSelect selectBody) {
+
+		if (CollectionUtils.isEmpty(selectBody.getSelectItems())) {
+			return new HashSet<>();
+		}
+
+		return selectBody.getSelectItems().stream() //
+				.filter(SelectExpressionItem.class::isInstance) //
+				.map(item -> ((SelectExpressionItem) item).getAlias()) //
+				.filter(Objects::nonNull) //
+				.map(Alias::getName) //
+				.collect(Collectors.toSet());
+	}
+
+	/**
+	 * Returns the aliases used for {@code join}s.
+	 *
+	 * @param query a query string to extract the aliases of joins from. Must not be {@literal null}.
+	 * @return a {@literal Set} of aliases used in the query. Guaranteed to be not {@literal null}.
+	 */
+	private Set<String> getJoinAliases(String query) {
+		return getJoinAliases((PlainSelect) parseSelectStatement(query).getSelectBody());
+	}
+
+	/**
+	 * Returns the aliases used for {@code join}s.
+	 *
+	 * @param selectBody the selection body to extract the aliases of joins from. Must not be {@literal null}.
+	 * @return a {@literal Set} of aliases used in the query. Guaranteed to be not {@literal null}.
+	 */
+	private Set<String> getJoinAliases(PlainSelect selectBody) {
+
+		if (CollectionUtils.isEmpty(selectBody.getJoins())) {
+			return new HashSet<>();
+		}
+
+		return selectBody.getJoins().stream() //
+				.map(join -> join.getRightItem().getAlias()) //
+				.filter(Objects::nonNull) //
+				.map(Alias::getName) //
+				.collect(Collectors.toSet());
+	}
+
+	/**
+	 * Returns the order clause for the given {@link Sort.Order}. Will prefix the clause with the given alias if the
+	 * referenced property refers to a join alias, i.e. starts with {@code $alias.}.
+	 *
+	 * @param joinAliases the join aliases of the original query. Must not be {@literal null}.
+	 * @param alias the alias for the root entity. May be {@literal null}.
+	 * @param order the order object to build the clause for. Must not be {@literal null}.
+	 * @return a {@link OrderByElement} containing an order clause. Guaranteed to be not {@literal null}.
+	 */
+	private OrderByElement getOrderClause(final Set<String> joinAliases, final Set<String> selectionAliases,
+			final String alias, final Sort.Order order) {
+
+		final OrderByElement orderByElement = new OrderByElement();
+		orderByElement.setAsc(order.getDirection().isAscending());
+		orderByElement.setAscDescPresent(true);
+
+		final String property = order.getProperty();
+
+		checkSortExpression(order);
+
+		if (selectionAliases.contains(property)) {
+			Expression orderExpression = order.isIgnoreCase() ? getJSqlLower(property) : new Column(property);
+
+			orderByElement.setExpression(orderExpression);
+			return orderByElement;
+		}
+
+		boolean qualifyReference = joinAliases //
+				.parallelStream() //
+				.map(joinAlias -> joinAlias.concat(".")) //
+				.noneMatch(property::startsWith);
+
+		boolean functionIndicator = property.contains("(");
+
+		String reference = qualifyReference && !functionIndicator && StringUtils.hasText(alias)
+				? String.format("%s.%s", alias, property)
+				: property;
+		Expression orderExpression = order.isIgnoreCase() ? getJSqlLower(reference) : new Column(reference);
+		orderByElement.setExpression(orderExpression);
+		return orderByElement;
 	}
 
 	@Override
 	public String detectAlias() {
-		return JSqlParserQueryUtils.detectAlias(this.query.getQueryString());
+		return detectAlias(this.query.getQueryString());
+	}
+
+	/**
+	 * Resolves the alias for the entity to be retrieved from the given JPA query. Note that you only provide valid Query
+	 * strings. Things such as <code>from User u</code> will throw an {@link IllegalArgumentException}.
+	 *
+	 * @param query must not be {@literal null}.
+	 * @return Might return {@literal null}.
+	 */
+	private 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();
 	}
 
 	@Override
 	public String createCountQueryFor(String countProjection) {
-		return JSqlParserQueryUtils.createCountQueryFor(this.query.getQueryString(), countProjection);
+
+		Assert.hasText(this.query.getQueryString(), "OriginalQuery must not be null or empty!");
+
+		Select selectStatement = parseSelectStatement(this.query.getQueryString());
+		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();
+
 	}
 
 	@Override
 	public String getProjection() {
-		return JSqlParserQueryUtils.getProjection(this.query.getQueryString());
+		Assert.hasText(query.getQueryString(), "Query must not be null or empty!");
+
+		Select selectStatement = parseSelectStatement(query.getQueryString());
+		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+
+		return selectBody.getSelectItems() //
+				.stream() //
+				.map(Object::toString) //
+				.collect(Collectors.joining(", ")).trim();
+	}
+
+	@Override
+	public Set<String> getJoinAliases() {
+		return this.getJoinAliases(this.query.getQueryString());
+	}
+
+	/**
+	 * Parses a query string with JSqlParser.
+	 *
+	 * @param query the query to parse
+	 * @return the parsed query
+	 */
+	private static Select parseSelectStatement(String query) {
+		try {
+			return (Select) CCJSqlParserUtil.parse(query);
+		} catch (JSQLParserException e) {
+			throw new IllegalArgumentException("The query you provided is not a valid SQL Query!", e);
+		}
+	}
+
+	/**
+	 * 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 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;
 	}
 
 	@Override