Differentiate between JPQL and native queries in count query derivation

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/DefaultQueryEnhancerUnitTests.java

@@ -15,6 +15,9 @@
  */
 package org.springframework.data.jpa.repository.query;
 
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
 /**
  * TCK Tests for {@link DefaultQueryEnhancer}.
  *
@@ -27,4 +30,8 @@ QueryEnhancer createQueryEnhancer(DeclaredQuery declaredQuery) {
 		return new DefaultQueryEnhancer(declaredQuery);
 	}
 
+	@Override
+	@Test // GH-2511, GH-2773
+	@Disabled("Not properly supported by QueryUtils")
+	void shouldDeriveNativeCountQueryWithVariable(String query, String expected) {}
 }

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancerUnitTests.java

@@ -15,10 +15,23 @@
  */
 package org.springframework.data.jpa.repository.query;
 
+import static org.assertj.core.api.Assertions.*;
+import static org.assertj.core.api.Assumptions.*;
+
+import java.util.stream.Stream;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.springframework.data.domain.Sort;
+
 /**
  * TCK Tests for {@link JSqlParserQueryEnhancer}.
  *
  * @author Mark Paluch
+ * @author Diego Krupitza
+ * @author Geoffrey Deremetz
  */
 public class JSqlParserQueryEnhancerUnitTests extends QueryEnhancerTckTests {
 
@@ -27,4 +40,181 @@ QueryEnhancer createQueryEnhancer(DeclaredQuery declaredQuery) {
 		return new JSqlParserQueryEnhancer(declaredQuery);
 	}
 
+	@Override
+	@ParameterizedTest // GH-2773
+	@MethodSource("jpqlCountQueries")
+	void shouldDeriveJpqlCountQuery(String query, String expected) {
+
+		assumeThat(query).as("JSQLParser does not support simple JPQL syntax").doesNotStartWithIgnoringCase("FROM");
+
+		assumeThat(query).as("JSQLParser does not support constructor JPQL syntax").doesNotContain(" new ");
+
+		super.shouldDeriveJpqlCountQuery(query, expected);
+	}
+
+	@Test
+	// GH-2578
+	void setOperationListWorks() {
+
+		String setQuery = "select SOME_COLUMN from SOME_TABLE where REPORTING_DATE = :REPORTING_DATE  \n" //
+				+ "except \n" //
+				+ "select SOME_COLUMN from SOME_OTHER_TABLE where REPORTING_DATE = :REPORTING_DATE";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isNullOrEmpty();
+		assertThat(stringQuery.getProjection()).isEqualToIgnoringCase("SOME_COLUMN");
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(setQuery);
+		assertThat(queryEnhancer.applySorting(Sort.by("SOME_COLUMN"))).endsWith("ORDER BY SOME_COLUMN ASC");
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isNullOrEmpty();
+		assertThat(queryEnhancer.getProjection()).isEqualToIgnoringCase("SOME_COLUMN");
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@Test // GH-2578
+	void complexSetOperationListWorks() {
+
+		String setQuery = "select SOME_COLUMN from SOME_TABLE where REPORTING_DATE = :REPORTING_DATE  \n" //
+				+ "except \n" //
+				+ "select SOME_COLUMN from SOME_OTHER_TABLE where REPORTING_DATE = :REPORTING_DATE \n" //
+				+ "union select SOME_COLUMN from SOME_OTHER_OTHER_TABLE";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isNullOrEmpty();
+		assertThat(stringQuery.getProjection()).isEqualToIgnoringCase("SOME_COLUMN");
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(setQuery);
+		assertThat(queryEnhancer.applySorting(Sort.by("SOME_COLUMN").ascending())).endsWith("ORDER BY SOME_COLUMN ASC");
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isNullOrEmpty();
+		assertThat(queryEnhancer.getProjection()).isEqualToIgnoringCase("SOME_COLUMN");
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@Test // GH-2578
+	void deeplyNestedcomplexSetOperationListWorks() {
+
+		String setQuery = "SELECT CustomerID FROM (\n" //
+				+ "\t\t\tselect * from Customers\n" //
+				+ "\t\t\texcept\n"//
+				+ "\t\t\tselect * from Customers where country = 'Austria'\n"//
+				+ "\t)\n" //
+				+ "\texcept\n"//
+				+ "\tselect CustomerID  from customers where country = 'Germany'\n"//
+				+ "\t;";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isNullOrEmpty();
+		assertThat(stringQuery.getProjection()).isEqualToIgnoringCase("CustomerID");
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(setQuery);
+		assertThat(queryEnhancer.applySorting(Sort.by("CustomerID").descending())).endsWith("ORDER BY CustomerID DESC");
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isNullOrEmpty();
+		assertThat(queryEnhancer.getProjection()).isEqualToIgnoringCase("CustomerID");
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@Test // GH-2578
+	void valuesStatementsWorks() {
+
+		String setQuery = "VALUES (1, 2, 'test')";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isNullOrEmpty();
+		assertThat(stringQuery.getProjection()).isNullOrEmpty();
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(setQuery);
+		assertThat(queryEnhancer.applySorting(Sort.by("CustomerID").descending())).isEqualTo(setQuery);
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isNullOrEmpty();
+		assertThat(queryEnhancer.getProjection()).isNullOrEmpty();
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@Test // GH-2578
+	void withStatementsWorks() {
+
+		String setQuery = "with sample_data(day, value) as (values ((0, 13), (1, 12), (2, 15), (3, 4), (4, 8), (5, 16))) \n"
+				+ "select day, value from sample_data as a";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isEqualToIgnoringCase("a");
+		assertThat(stringQuery.getProjection()).isEqualToIgnoringCase("day, value");
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(
+				"with sample_data (day, value) AS (VALUES ((0, 13), (1, 12), (2, 15), (3, 4), (4, 8), (5, 16)))\n"
+						+ "SELECT count(1) FROM sample_data AS a");
+		assertThat(queryEnhancer.applySorting(Sort.by("day").descending())).endsWith("ORDER BY a.day DESC");
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isEqualToIgnoringCase("a");
+		assertThat(queryEnhancer.getProjection()).isEqualToIgnoringCase("day, value");
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@Test // GH-2578
+	void multipleWithStatementsWorks() {
+
+		String setQuery = "with sample_data(day, value) as (values ((0, 13), (1, 12), (2, 15), (3, 4), (4, 8), (5, 16))), test2 as (values (1,2,3)) \n"
+				+ "select day, value from sample_data as a";
+
+		StringQuery stringQuery = new StringQuery(setQuery, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(stringQuery.getAlias()).isEqualToIgnoringCase("a");
+		assertThat(stringQuery.getProjection()).isEqualToIgnoringCase("day, value");
+		assertThat(stringQuery.hasConstructorExpression()).isFalse();
+
+		assertThat(queryEnhancer.createCountQueryFor()).isEqualToIgnoringCase(
+				"with sample_data (day, value) AS (VALUES ((0, 13), (1, 12), (2, 15), (3, 4), (4, 8), (5, 16))),test2 AS (VALUES (1, 2, 3))\n"
+						+ "SELECT count(1) FROM sample_data AS a");
+		assertThat(queryEnhancer.applySorting(Sort.by("day").descending())).endsWith("ORDER BY a.day DESC");
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isEqualToIgnoringCase("a");
+		assertThat(queryEnhancer.getProjection()).isEqualToIgnoringCase("day, value");
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	@ParameterizedTest // GH-2641
+	@MethodSource("mergeStatementWorksSource")
+	void mergeStatementWorksWithJSqlParser(String query, String alias) {
+
+		StringQuery stringQuery = new StringQuery(query, true);
+		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
+
+		assertThat(queryEnhancer.detectAlias()).isEqualTo(alias);
+		assertThat(QueryUtils.detectAlias(query)).isNull();
+
+		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
+		assertThat(queryEnhancer.detectAlias()).isEqualTo(alias);
+		assertThat(queryEnhancer.getProjection()).isEmpty();
+		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
+	}
+
+	static Stream<Arguments> mergeStatementWorksSource() {
+
+		return Stream.of( //
+				Arguments.of(
+						"merge into a using (select id, value from b) query on (a.id = query.id) when matched then update set a.value = value",
+						"query"),
+				Arguments.of(
+						"merge into a using (select id2, value from b) on (id = id2) when matched then update set a.value = value",
+						null));
+	}
 }

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/QueryEnhancerTckTests.java

@@ -19,6 +19,7 @@
 
 import java.util.stream.Stream;
 
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -36,9 +37,13 @@ void shouldDeriveNativeCountQuery(String query, String expected) {
 
 		DeclaredQuery declaredQuery = DeclaredQuery.of(query, true);
 		QueryEnhancer enhancer = createQueryEnhancer(declaredQuery);
-		String countQueryFor = enhancer.createCountQueryFor(null);
+		String countQueryFor = enhancer.createCountQueryFor();
 
-		assertThat(countQueryFor).isEqualToIgnoringCase(expected);
+		// lenient cleanup to allow for rendering variance
+		String sanitized = countQueryFor.replaceAll("\r", " ").replaceAll("\n", " ").replaceAll(" {2}", " ")
+				.replaceAll(" {2}", " ").trim();
+
+		assertThat(sanitized).isEqualToIgnoringCase(expected);
 	}
 
 	static Stream<Arguments> nativeCountQueries() {
@@ -53,7 +58,58 @@ static Stream<Arguments> nativeCountQueries() {
 
 				Arguments.of( //
 						"SELECT DISTINCT name FROM table_name some_alias", //
-						"select count(DISTINCT name) FROM table_name some_alias"));
+						"select count(DISTINCT name) FROM table_name some_alias"), //
+
+				Arguments.of( //
+						"select distinct u from User u where u.foo = ?", //
+						"select count(distinct u) from User u where u.foo = ?"),
+
+				Arguments.of( //
+						"select u from User as u", //
+						"select count(u) from User as u"),
+
+				Arguments.of( //
+						"SELECT u FROM User u where u.foo.bar = ?", //
+						"select count(u) FROM User u where u.foo.bar = ?"),
+
+				Arguments.of( //
+						"select p.lastname,p.firstname from Person p", //
+						"select count(1) from Person p"),
+
+				// whitespace quirks
+				Arguments.of( //
+						"""
+								select user.age,
+								  user.name
+								  from User user
+								  where user.age = 18
+								  order
+								by
+								user.name
+								\s""", //
+						"select count(1) from User user where user.age = 18"),
+
+				Arguments.of( //
+						"select * from User user\n" + //
+								"  where user.age = 18\n" + //
+								"  order by user.name\n ", //
+						"select count(1) from User user where user.age = 18"),
+
+				Arguments.of( //
+						"SELECT DISTINCT entity1\nFROM Entity1 entity1\nLEFT JOIN Entity2 entity2 ON entity1.key = entity2.key", //
+						"select count(DISTINCT entity1) FROM Entity1 entity1 LEFT JOIN Entity2 entity2 ON entity1.key = entity2.key"),
+
+				Arguments.of( //
+						"SELECT DISTINCT entity1\nFROM Entity1 entity1 LEFT JOIN Entity2 entity2 ON entity1.key = entity2.key", //
+						"select count(DISTINCT entity1) FROM Entity1 entity1 LEFT JOIN Entity2 entity2 ON entity1.key = entity2.key"),
+
+				Arguments.of( //
+						"SELECT DISTINCT entity1\nFROM Entity1 entity1 LEFT JOIN Entity2 entity2 ON entity1.key = entity2.key\nwhere entity1.id = 1799", //
+						"select count(DISTINCT entity1) FROM Entity1 entity1 LEFT JOIN Entity2 entity2 ON entity1.key = entity2.key where entity1.id = 1799"),
+
+				Arguments.of( //
+						"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"));
 	}
 
 	@ParameterizedTest // GH-2773
@@ -79,7 +135,70 @@ static Stream<Arguments> jpqlCountQueries() {
 
 				Arguments.of( //
 						"SELECT DISTINCT name FROM table_name some_alias", //
-						"select count(DISTINCT name) FROM table_name some_alias"));
+						"select count(DISTINCT name) FROM table_name some_alias"),
+
+				Arguments.of( //
+						"select distinct new User(u.name) from User u where u.foo = ?", //
+						"select count(distinct u) from User u where u.foo = ?"),
+
+				Arguments.of( //
+						"FROM User u WHERE u.foo.bar = ?", //
+						"select count(u) FROM User u WHERE u.foo.bar = ?"),
+
+				Arguments.of( //
+						"from User u", //
+						"select count(u) FROM User u"),
+
+				Arguments.of( //
+						"select u from User as u", //
+						"select count(u) from User as u"),
+
+				Arguments.of( //
+						"select p.lastname,p.firstname from Person p", //
+						"select count(p) from Person p"),
+
+				Arguments.of( //
+						"select a.b from A a", //
+						"select count(a.b) from A a"),
+
+				Arguments.of( //
+						"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"));
+	}
+
+	@ParameterizedTest // GH-2511, GH-2773
+	@MethodSource("nativeQueriesWithVariables")
+	void shouldDeriveNativeCountQueryWithVariable(String query, String expected) {
+
+		DeclaredQuery declaredQuery = DeclaredQuery.of(query, true);
+		QueryEnhancer enhancer = createQueryEnhancer(declaredQuery);
+		String countQueryFor = enhancer.createCountQueryFor();
+
+		assertThat(countQueryFor).isEqualToIgnoringCase(expected);
+	}
+
+	static Stream<Arguments> nativeQueriesWithVariables() {
+
+		return Stream.of(Arguments.of( //
+				"SELECT * FROM User WHERE created_at > $1", //
+				"SELECT count(1) FROM User WHERE created_at > $1"), //
+
+				Arguments.of( //
+						"SELECT * FROM (select * from test) ", //
+						"SELECT count(1) FROM (SELECT * FROM test)"), //
+
+				Arguments.of( //
+						"SELECT * FROM (select * from test) as test", //
+						"SELECT count(1) FROM (SELECT * FROM test) AS test"));
+	}
+
+	@Test
+	// DATAJPA-1696
+	void findProjectionClauseWithIncludedFrom() {
+
+		StringQuery query = new StringQuery("select x, frommage, y from t", true);
+
+		assertThat(createQueryEnhancer(query).getProjection()).isEqualTo("x, frommage, y");
 	}
 
 	abstract QueryEnhancer createQueryEnhancer(DeclaredQuery declaredQuery);