fix: fix query generation for three or more keys

Signed-off-by: Artemiy Chereshnevvv <[email protected]>

Author: chrshnv

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/query/StatementFactory.java

@@ -249,39 +249,42 @@ Criteria applyScrollCriteria(Criteria criteria, @Nullable ScrollPosition positio
 			List<String> columns = new ArrayList<>(keys.keySet());
 			List<Object> values = new ArrayList<>(keys.values());
 
+			if (columns.isEmpty() || values.isEmpty())
+				return criteria;
+
 			Map<String, Sort.Direction> directions =
 				sort.stream().collect(Collectors.toMap(Sort.Order::getProperty, Sort.Order::getDirection));
 
-			String primary = columns.get(0);
-			Object primaryValue = values.get(0);
+			Sort.Direction dir = directions.getOrDefault(columns.get(0), Sort.DEFAULT_DIRECTION);
 
-			Criteria primaryCompare;
-			Sort.Direction dir = directions.getOrDefault(primary, Sort.DEFAULT_DIRECTION);
+			Criteria scroll = buildKeysetCriteria(columns, values, keyset.scrollsForward(), dir);
 
-			if (keyset.scrollsForward() ^ dir.isDescending()) {
-				primaryCompare = columns.size() != 1 ? Criteria.where(primary).greaterThanOrEquals(primaryValue) : Criteria.where(primary).greaterThan(primaryValue);
-			} else {
-				primaryCompare = columns.size() != 1 ? Criteria.where(primary).lessThanOrEquals(primaryValue) : Criteria.where(primary).lessThan(primaryValue);
-			}
+			return criteria.and(scroll);
+		}
 
-			if (columns.size() == 1)
-				return criteria.and(primaryCompare);
+		Criteria buildKeysetCriteria(List<String> columns, List<Object> values, boolean isForward, Sort.Direction dir) {
+			if (columns.isEmpty())
+				return Criteria.empty();
 
-			Criteria tupleCompare = null;
-			for (int i = 0; i < columns.size(); i++) {
-				Criteria orCriteria;
-				String col = columns.get(i);
-				Object val = values.get(i);
+			String column = columns.get(0);
+			Object value = values.get(0);
 
-				Sort.Direction colDir = directions.getOrDefault(col, dir);
-				boolean asc = keyset.scrollsForward() ? colDir.isAscending() : colDir.isDescending();
+			boolean isAscending = isForward ^ dir.isDescending();
 
-				orCriteria = asc ? Criteria.where(col).greaterThan(val) : Criteria.where(col).lessThan(val);
+			Criteria gte = isAscending
+				? Criteria.where(column).greaterThanOrEquals(value)
+				: Criteria.where(column).lessThanOrEquals(value);
 
-				tupleCompare = (tupleCompare == null) ? orCriteria : tupleCompare.or(orCriteria);
-			}
+			Criteria gt = isAscending
+				? Criteria.where(column).greaterThan(value)
+				: Criteria.where(column).lessThan(value);
+
+			if (columns.size() == 1)
+				return gt;
+
+			Criteria nested = buildKeysetCriteria(columns.subList(1, columns.size()), values.subList(1, values.size()), isForward, dir);
 
-			return criteria.and(primaryCompare).and(tupleCompare);
+			return gte.and(gt.or(nested));
 		}
 
 		SelectBuilder.SelectOrdered applyOrderBy(Sort sort, RelationalPersistentEntity<?> entity, Table table,

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/query/PartTreeJdbcQueryUnitTests.java

@@ -744,13 +744,37 @@ void createsQueryForCountProjection() throws Exception {
 	void createQueryForKeysetTwoKeys() throws Exception {
 		JdbcQueryMethod queryMethod = getQueryMethod("findFirst2ByOrderByIdIdAsc", ScrollPosition.class);
 		PartTreeJdbcQuery jdbcQuery = createQuery(queryMethod);
-		ParametrizedQuery query = jdbcQuery.createQuery((getAccessor(queryMethod, new Object[]{ScrollPosition.of(new LinkedHashMap<>(Map.of("idId", 1, "firstName", "John")), ScrollPosition.Direction.FORWARD)})), returnedType);
+		Map<String, Object> values = new LinkedHashMap<>() {{
+			put("idId", 1);
+			put("firstName", "John");
+		}};
+
+		ParametrizedQuery query = jdbcQuery.createQuery((getAccessor(queryMethod, new Object[]{ScrollPosition.of(values, ScrollPosition.Direction.FORWARD)})), returnedType);
 
 		QueryAssert.assertThat(query).containsQuotedAliasedColumns(columns)
-			.contains(" WHERE (" + TABLE + ".\"ID\" >= :id) AND (" + TABLE + ".\"ID\" > :id1 OR (" + TABLE + ".\"FIRST_NAME\" > :first_name").hasBindValue("id", 1L)
+			.contains(" WHERE (" + TABLE + ".\"ID\" >= :id AND (" + TABLE + ".\"ID\" > :id1 OR (" + TABLE + ".\"FIRST_NAME\" > :first_name").hasBindValue("id", 1L)
 			.hasBindValue("first_name", "John");
 	}
 
+	@Test // PR-2149
+	void createQueryForKeysetThreeKeys() throws Exception {
+		JdbcQueryMethod queryMethod = getQueryMethod("findFirst2ByOrderByIdIdAsc", ScrollPosition.class);
+		PartTreeJdbcQuery jdbcQuery = createQuery(queryMethod);
+		Map<String, Object> values = new LinkedHashMap<>() {{
+			put("idId", 1);
+			put("firstName", "John");
+			put("age", 18);
+		}};
+
+		ParametrizedQuery query = jdbcQuery.createQuery((getAccessor(queryMethod, new Object[]{ScrollPosition.of(values, ScrollPosition.Direction.FORWARD)})), returnedType);
+
+		QueryAssert.assertThat(query).containsQuotedAliasedColumns(columns)
+			.contains("WHERE (" + TABLE + ".\"ID\" >= :id AND (" + TABLE + ".\"ID\" > :id1 OR (" + TABLE + ".\"FIRST_NAME\" >= :first_name AND (" + TABLE + ".\"FIRST_NAME\" > :first_name1 OR (" + TABLE + ".\"AGE\" > :age)")
+			.hasBindValue("id", 1L)
+			.hasBindValue("first_name", "John")
+			.hasBindValue("age", 18);
+	}
+
 	private PartTreeJdbcQuery createQuery(JdbcQueryMethod queryMethod) {
 		return new PartTreeJdbcQuery(mappingContext, queryMethod, JdbcH2Dialect.INSTANCE, converter,
 			mock(NamedParameterJdbcOperations.class), mock(RowMapper.class));