Prevent sorting and count queries for non-SELECT statements.

Add statement type detection to QueryEnhancer implementations to validate
operations. QueryEnhancer now throws IllegalStateException when attempting
to create count queries or apply sorting to INSERT, UPDATE, DELETE, or MERGE
statements, as these operations are only valid for SELECT queries.

Signed-off-by: kssumin <ksoomin25@gmail.com>
Original pull request: #4052
Closes #2856

Author: kssumin

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

@@ -27,7 +27,7 @@ class DefaultQueryEnhancer implements QueryEnhancer {
 
 	private final QueryProvider query;
 	private final boolean hasConstructorExpression;
-	private final @Nullable  String alias;
+	private final @Nullable String alias;
 	private final String projection;
 
 	public DefaultQueryEnhancer(QueryProvider query) {

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

@@ -28,6 +28,7 @@
  *
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author kssumin
  */
 @SuppressWarnings("UnreachableCode")
 class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIntrospector<QueryInformation> {
@@ -38,11 +39,36 @@ class EqlQueryIntrospector extends EqlBaseVisitor<Void> implements ParsedQueryIn
 	private @Nullable List<QueryToken> projection;
 	private boolean projectionProcessed;
 	private boolean hasConstructorExpression = false;
+	private QueryInformation.StatementType statementType = QueryInformation.StatementType.SELECT;
 
 	@Override
 	public QueryInformation getParsedQueryInformation() {
 		return new QueryInformation(primaryFromAlias, projection == null ? Collections.emptyList() : projection,
-				hasConstructorExpression);
+				hasConstructorExpression, statementType);
+	}
+
+	@Override
+	public Void visitSelectQuery(EqlParser.SelectQueryContext ctx) {
+		statementType = QueryInformation.StatementType.SELECT;
+		return super.visitSelectQuery(ctx);
+	}
+
+	@Override
+	public Void visitFromQuery(EqlParser.FromQueryContext ctx) {
+		statementType = QueryInformation.StatementType.SELECT;
+		return super.visitFromQuery(ctx);
+	}
+
+	@Override
+	public Void visitUpdate_statement(EqlParser.Update_statementContext ctx) {
+		statementType = QueryInformation.StatementType.UPDATE;
+		return super.visitUpdate_statement(ctx);
+	}
+
+	@Override
+	public Void visitDelete_statement(EqlParser.Delete_statementContext ctx) {
+		statementType = QueryInformation.StatementType.DELETE;
+		return super.visitDelete_statement(ctx);
 	}
 
 	@Override

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

@@ -24,28 +24,35 @@
  *
  * @author Mark Paluch
  * @author Oscar Fanchin
+ * @author kssumin
  * @since 3.5
  */
 class HibernateQueryInformation extends QueryInformation {
 
 	private final boolean hasCte;
-	
+
 	private final boolean hasFromFunction;
-	
 
 	public HibernateQueryInformation(@Nullable String alias, List<QueryToken> projection,
-			boolean hasConstructorExpression, boolean hasCte,boolean hasFromFunction) {
+			boolean hasConstructorExpression, boolean hasCte, boolean hasFromFunction) {
 		super(alias, projection, hasConstructorExpression);
 		this.hasCte = hasCte;
 		this.hasFromFunction = hasFromFunction;
 	}
 
+	public HibernateQueryInformation(@Nullable String alias, List<QueryToken> projection,
+			boolean hasConstructorExpression, StatementType statementType, boolean hasCte, boolean hasFromFunction) {
+		super(alias, projection, hasConstructorExpression, statementType);
+		this.hasCte = hasCte;
+		this.hasFromFunction = hasFromFunction;
+	}
+
 	public boolean hasCte() {
 		return hasCte;
 	}
-	
+
 	public boolean hasFromFunction() {
 		return hasFromFunction;
 	}
-	
+
 }

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

@@ -30,6 +30,7 @@
  *
  * @author Mark Paluch
  * @author Oscar Fanchin
+ * @author kssumin
  */
 @SuppressWarnings({ "UnreachableCode", "ConstantValue" })
 class HqlQueryIntrospector extends HqlBaseVisitor<Void> implements ParsedQueryIntrospector<HibernateQueryInformation> {
@@ -42,11 +43,36 @@ class HqlQueryIntrospector extends HqlBaseVisitor<Void> implements ParsedQueryIn
 	private boolean hasConstructorExpression = false;
 	private boolean hasCte = false;
 	private boolean hasFromFunction = false;
+	private QueryInformation.StatementType statementType = QueryInformation.StatementType.SELECT;
 
 	@Override
 	public HibernateQueryInformation getParsedQueryInformation() {
 		return new HibernateQueryInformation(primaryFromAlias, projection == null ? Collections.emptyList() : projection,
-				hasConstructorExpression, hasCte, hasFromFunction);
+				hasConstructorExpression, statementType, hasCte, hasFromFunction);
+	}
+
+	@Override
+	public Void visitSelectStatement(HqlParser.SelectStatementContext ctx) {
+		statementType = QueryInformation.StatementType.SELECT;
+		return super.visitSelectStatement(ctx);
+	}
+
+	@Override
+	public Void visitInsertStatement(HqlParser.InsertStatementContext ctx) {
+		statementType = QueryInformation.StatementType.INSERT;
+		return super.visitInsertStatement(ctx);
+	}
+
+	@Override
+	public Void visitUpdateStatement(HqlParser.UpdateStatementContext ctx) {
+		statementType = QueryInformation.StatementType.UPDATE;
+		return super.visitUpdateStatement(ctx);
+	}
+
+	@Override
+	public Void visitDeleteStatement(HqlParser.DeleteStatementContext ctx) {
+		statementType = QueryInformation.StatementType.DELETE;
+		return super.visitDeleteStatement(ctx);
 	}
 
 	@Override

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

@@ -69,6 +69,7 @@
  * @author Yanming Zhou
  * @author Christoph Strobl
  * @author Diego Pedregal
+ * @author kssumin
  * @since 2.7.0
  */
 public class JSqlParserQueryEnhancer implements QueryEnhancer {
@@ -343,7 +344,16 @@ private String doApplySorting(Sort sort, @Nullable String alias) {
 		String queryString = query.getQueryString();
 		Assert.hasText(queryString, "Query must not be null or empty");
 
-		if (this.parsedType != ParsedType.SELECT || sort.isUnsorted()) {
+		if (this.parsedType != ParsedType.SELECT) {
+			if (!sort.isUnsorted()) {
+				throw new IllegalStateException(String.format(
+						"Cannot apply sorting to %s statement. Sorting is only supported for SELECT statements.",
+						this.parsedType));
+			}
+			return queryString;
+		}
+
+		if (sort.isUnsorted()) {
 			return queryString;
 		}
 
@@ -383,7 +393,9 @@ private String applySorting(@Nullable Select selectStatement, Sort sort, @Nullab
 	public String createCountQueryFor(@Nullable String countProjection) {
 
 		if (this.parsedType != ParsedType.SELECT) {
-			return this.query.getQueryString();
+			throw new IllegalStateException(String.format(
+					"Cannot derive count query for %s statement. Count queries are only supported for SELECT statements.",
+					this.parsedType));
 		}
 
 		Assert.hasText(this.query.getQueryString(), "OriginalQuery must not be null or empty");

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

@@ -41,6 +41,7 @@
  *
  * @author Greg Turnquist
  * @author Mark Paluch
+ * @author kssumin
  * @since 3.1
  * @see JpqlQueryParser
  * @see HqlQueryParser
@@ -220,6 +221,13 @@ public DeclaredQuery getQuery() {
 
 	@Override
 	public String rewrite(QueryRewriteInformation rewriteInformation) {
+
+		if (!queryInformation.isSelectStatement() && !rewriteInformation.getSort().isUnsorted()) {
+			throw new IllegalStateException(String.format(
+					"Cannot apply sorting to %s statement. Sorting is only supported for SELECT statements.",
+					queryInformation.getStatementType()));
+		}
+
 		return QueryRenderer.TokenRenderer.render(
 				sortFunction.apply(rewriteInformation.getSort(), this.queryInformation, rewriteInformation.getReturnedType())
 						.visit(context));
@@ -232,6 +240,13 @@ public String rewrite(QueryRewriteInformation rewriteInformation) {
 	 */
 	@Override
 	public String createCountQueryFor(@Nullable String countProjection) {
+
+		if (!queryInformation.isSelectStatement()) {
+			throw new IllegalStateException(String.format(
+					"Cannot derive count query for %s statement. Count queries are only supported for SELECT statements.",
+					queryInformation.getStatementType()));
+		}
+
 		return QueryRenderer.TokenRenderer
 				.render(countQueryFunction.apply(countProjection, this.queryInformation).visit(context));
 	}

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

@@ -28,6 +28,7 @@
  *
  * @author Mark Paluch
  * @author Christoph Strobl
+ * @author kssumin
  */
 @SuppressWarnings({ "UnreachableCode", "ConstantValue" })
 class JpqlQueryIntrospector extends JpqlBaseVisitor<Void> implements ParsedQueryIntrospector<QueryInformation> {
@@ -38,11 +39,36 @@ class JpqlQueryIntrospector extends JpqlBaseVisitor<Void> implements ParsedQuery
 	private @Nullable List<QueryToken> projection;
 	private boolean projectionProcessed;
 	private boolean hasConstructorExpression = false;
+	private QueryInformation.StatementType statementType = QueryInformation.StatementType.SELECT;
 
 	@Override
 	public QueryInformation getParsedQueryInformation() {
 		return new QueryInformation(primaryFromAlias, projection == null ? Collections.emptyList() : projection,
-				hasConstructorExpression);
+				hasConstructorExpression, statementType);
+	}
+
+	@Override
+	public Void visitSelectQuery(JpqlParser.SelectQueryContext ctx) {
+		statementType = QueryInformation.StatementType.SELECT;
+		return super.visitSelectQuery(ctx);
+	}
+
+	@Override
+	public Void visitFromQuery(JpqlParser.FromQueryContext ctx) {
+		statementType = QueryInformation.StatementType.SELECT;
+		return super.visitFromQuery(ctx);
+	}
+
+	@Override
+	public Void visitUpdate_statement(JpqlParser.Update_statementContext ctx) {
+		statementType = QueryInformation.StatementType.UPDATE;
+		return super.visitUpdate_statement(ctx);
+	}
+
+	@Override
+	public Void visitDelete_statement(JpqlParser.Delete_statementContext ctx) {
+		statementType = QueryInformation.StatementType.DELETE;
+		return super.visitDelete_statement(ctx);
 	}
 
 	@Override

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

@@ -23,6 +23,7 @@
  * Value object capturing introspection details of a parsed query.
  *
  * @author Mark Paluch
+ * @author kssumin
  * @since 3.5
  */
 class QueryInformation {
@@ -33,10 +34,18 @@ class QueryInformation {
 
 	private final boolean hasConstructorExpression;
 
+	private final StatementType statementType;
+
 	QueryInformation(@Nullable String alias, List<QueryToken> projection, boolean hasConstructorExpression) {
+		this(alias, projection, hasConstructorExpression, StatementType.SELECT);
+	}
+
+	QueryInformation(@Nullable String alias, List<QueryToken> projection, boolean hasConstructorExpression,
+			StatementType statementType) {
 		this.alias = alias;
 		this.projection = projection;
 		this.hasConstructorExpression = hasConstructorExpression;
+		this.statementType = statementType;
 	}
 
 	/**
@@ -61,4 +70,59 @@ public List<QueryToken> getProjection() {
 	public boolean hasConstructorExpression() {
 		return hasConstructorExpression;
 	}
+
+	/**
+	 * @return the statement type of the query.
+	 * @since 4.0
+	 */
+	public StatementType getStatementType() {
+		return statementType;
+	}
+
+	/**
+	 * @return {@code true} if the query is a SELECT statement.
+	 * @since 4.0
+	 */
+	public boolean isSelectStatement() {
+		return statementType == StatementType.SELECT;
+	}
+
+	/**
+	 * Enum representing the type of SQL/JPQL statement.
+	 *
+	 * @since 4.0
+	 */
+	enum StatementType {
+
+		/**
+		 * SELECT statement.
+		 */
+		SELECT,
+
+		/**
+		 * INSERT statement.
+		 */
+		INSERT,
+
+		/**
+		 * UPDATE statement.
+		 */
+		UPDATE,
+
+		/**
+		 * DELETE statement.
+		 */
+		DELETE,
+
+		/**
+		 * MERGE statement.
+		 */
+		MERGE,
+
+		/**
+		 * Other statement types.
+		 */
+		OTHER
+	}
+
 }