DATAES-176 - firstN repository queries.

Original PR: spring-projects#305

Author: sothawo

src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java

@@ -362,16 +362,19 @@ public <T> T queryForObject(GetQuery query, Class<T> clazz, GetResultMapper mapp
 
 	@Override
 	public <T> T queryForObject(CriteriaQuery query, Class<T> clazz) {
-		Page<T> page = queryForPage(query, clazz);
-		Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found " + page.getTotalElements() + " results");
-		return page.getTotalElements() > 0 ? page.getContent().get(0) : null;
+		return getObjectFromPage(queryForPage(query, clazz));
 	}
 
 	@Override
 	public <T> T queryForObject(StringQuery query, Class<T> clazz) {
-		Page<T> page = queryForPage(query, clazz);
-		Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found " + page.getTotalElements() + " results");
-		return page.getTotalElements() > 0 ? page.getContent().get(0) : null;
+		return getObjectFromPage(queryForPage(query, clazz));
+	}
+
+	@Nullable
+	private <T> T getObjectFromPage(Page<T> page) {
+		int contentSize = page.getContent().size();
+		Assert.isTrue(contentSize < 2, "Expected 1 but found " + contentSize + " results");
+		return contentSize > 0 ? page.getContent().get(0) : null;
 	}
 
 	@Override
@@ -492,19 +495,24 @@ public <T> Page<T> queryForPage(CriteriaQuery criteriaQuery, Class<T> clazz) {
 		QueryBuilder elasticsearchFilter = new CriteriaFilterProcessor()
 				.createFilterFromCriteria(criteriaQuery.getCriteria());
 		SearchRequest request = prepareSearch(criteriaQuery, clazz);
+		SearchSourceBuilder sourceBuilder = request.source();
 
 		if (elasticsearchQuery != null) {
-			request.source().query(elasticsearchQuery);
+			sourceBuilder.query(elasticsearchQuery);
 		} else {
-			request.source().query(QueryBuilders.matchAllQuery());
+			sourceBuilder.query(QueryBuilders.matchAllQuery());
+		}
+
+		if (criteriaQuery.isLimiting()) {
+			sourceBuilder.size(criteriaQuery.getMaxResults());
 		}
 
 		if (criteriaQuery.getMinScore() > 0) {
-			request.source().minScore(criteriaQuery.getMinScore());
+			sourceBuilder.minScore(criteriaQuery.getMinScore());
 		}
 
 		if (elasticsearchFilter != null)
-			request.source().postFilter(elasticsearchFilter);
+			sourceBuilder.postFilter(elasticsearchFilter);
 		if (logger.isDebugEnabled()) {
 			logger.debug("doSearch query:\n" + request.toString());
 		}

src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java

@@ -294,16 +294,19 @@ public <T> T queryForObject(GetQuery query, Class<T> clazz, GetResultMapper mapp
 
 	@Override
 	public <T> T queryForObject(CriteriaQuery query, Class<T> clazz) {
-		Page<T> page = queryForPage(query, clazz);
-		Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found " + page.getTotalElements() + " results");
-		return page.getTotalElements() > 0 ? page.getContent().get(0) : null;
+		return getObjectFromPage(queryForPage(query, clazz));
 	}
 
 	@Override
 	public <T> T queryForObject(StringQuery query, Class<T> clazz) {
-		Page<T> page = queryForPage(query, clazz);
-		Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found " + page.getTotalElements() + " results");
-		return page.getTotalElements() > 0 ? page.getContent().get(0) : null;
+		return getObjectFromPage(queryForPage(query, clazz));
+	}
+
+	@Nullable
+	private <T> T getObjectFromPage(Page<T> page) {
+		int contentSize = page.getContent().size();
+		Assert.isTrue(contentSize < 2, "Expected 1 but found " + contentSize + " results");
+		return contentSize > 0 ? page.getContent().get(0) : null;
 	}
 
 	@Override
@@ -423,6 +426,10 @@ public <T> Page<T> queryForPage(CriteriaQuery criteriaQuery, Class<T> clazz) {
 			searchRequestBuilder.setQuery(QueryBuilders.matchAllQuery());
 		}
 
+		if (criteriaQuery.isLimiting()) {
+			searchRequestBuilder.setSize(criteriaQuery.getMaxResults());
+		}
+
 		if (criteriaQuery.getMinScore() > 0) {
 			searchRequestBuilder.setMinScore(criteriaQuery.getMinScore());
 		}

src/main/java/org/springframework/data/elasticsearch/core/ReactiveElasticsearchTemplate.java

@@ -294,7 +294,12 @@ private Flux<SearchHit> doFind(Query query, ElasticsearchPersistentEntity<?> ent
 
 				request.source(searchSourceBuilder);
 				return doFind(prepareSearchRequest(request));
+			} else if (query.isLimiting()) {
+				searchSourceBuilder.from(0);
+				searchSourceBuilder.size(query.getMaxResults());
 
+				request.source(searchSourceBuilder);
+				return doFind(prepareSearchRequest(request));
 			} else {
 
 				request.source(searchSourceBuilder);
@@ -657,9 +662,7 @@ private QueryBuilder mappedQuery(Query query, ElasticsearchPersistentEntity<?> e
 			elasticsearchQuery = new WrapperQueryBuilder(((StringQuery) query).getSource());
 		} else if (query instanceof NativeSearchQuery) {
 			elasticsearchQuery = ((NativeSearchQuery) query).getQuery();
-		}
-
-		else {
+		} else {
 			throw new IllegalArgumentException(String.format("Unknown query type '%s'.", query.getClass()));
 		}
 

src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java

@@ -36,6 +36,7 @@
  * @author Alen Turkovic
  * @author Sascha Woo
  * @author Farid Azaza
+ * @author Peter-Josef Meisch
  */
 abstract class AbstractQuery implements Query {
 
@@ -52,6 +53,7 @@ abstract class AbstractQuery implements Query {
 	protected IndicesOptions indicesOptions;
 	protected boolean trackScores;
 	protected String preference;
+	protected Integer maxResults;
 
 	@Override
 	public Sort getSort() {
@@ -112,6 +114,7 @@ public SourceFilter getSourceFilter() {
 		return sourceFilter;
 	}
 
+	@Override
 	@SuppressWarnings("unchecked")
 	public final <T extends Query> T addSort(Sort sort) {
 		if (sort == null) {
@@ -127,6 +130,7 @@ public final <T extends Query> T addSort(Sort sort) {
 		return (T) this;
 	}
 
+	@Override
 	public float getMinScore() {
 		return minScore;
 	}
@@ -135,6 +139,7 @@ public void setMinScore(float minScore) {
 		this.minScore = minScore;
 	}
 
+	@Override
 	public Collection<String> getIds() {
 		return ids;
 	}
@@ -143,6 +148,7 @@ public void setIds(Collection<String> ids) {
 		this.ids = ids;
 	}
 
+	@Override
 	public String getRoute() {
 		return route;
 	}
@@ -155,10 +161,12 @@ public void setSearchType(SearchType searchType) {
 		this.searchType = searchType;
 	}
 
+	@Override
 	public SearchType getSearchType() {
 		return searchType;
 	}
 
+	@Override
 	public IndicesOptions getIndicesOptions() {
 		return indicesOptions;
 	}
@@ -178,7 +186,7 @@ public boolean getTrackScores() {
 
 	/**
 	 * Configures whether to track scores.
-	 * 
+	 *
 	 * @param trackScores
 	 * @since 3.1
 	 */
@@ -195,4 +203,18 @@ public String getPreference() {
 	public void setPreference(String preference) {
 		this.preference = preference;
 	}
+
+	@Override
+	public boolean isLimiting() {
+		return maxResults != null;
+	}
+
+	@Override
+	public Integer getMaxResults() {
+		return maxResults;
+	}
+
+	public void setMaxResults(Integer maxResults) {
+		this.maxResults = maxResults;
+	}
 }

src/main/java/org/springframework/data/elasticsearch/core/query/Query.java

@@ -35,6 +35,7 @@
  * @author Sascha Woo
  * @author Christoph Strobl
  * @author Farid Azaza
+ * @author Peter-Josef Meisch
  */
 public interface Query {
 
@@ -145,7 +146,7 @@ static Query findAll() {
 
 	/**
 	 * Get if scores will be computed and tracked, regardless of whether sorting on a field. Defaults to <tt>false</tt>.
-	 * 
+	 *
 	 * @return
 	 * @since 3.1
 	 */
@@ -194,4 +195,22 @@ static Query findAll() {
 	 * @since 3.2
 	 */
 	void setPreference(String preference);
+
+	/**
+	 * @return true if the query has a limit on the max number of results.
+	 * @since 4.0
+	 */
+	default boolean isLimiting() {
+		return false;
+	}
+
+	/**
+	 * return the max of results. Must not return null when {@link #isLimiting()} returns true.
+	 *
+	 * @since 4.0
+	 */
+	default Integer getMaxResults() {
+		return null;
+	}
+
 }

src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java

@@ -58,6 +58,10 @@ public Object execute(Object[] parameters) {
 		CriteriaQuery query = createQuery(accessor);
 		Assert.notNull(query, "unsupported query");
 
+		if (tree.isLimiting()) {
+			query.setMaxResults(tree.getMaxResults());
+		}
+
 		if (tree.isDelete()) {
 			Object result = countOrGetDocumentsForDelete(query, accessor);
 			elasticsearchOperations.delete(query, queryMethod.getEntityInformation().getJavaType());

src/main/java/org/springframework/data/elasticsearch/repository/query/ReactivePartTreeElasticsearchQuery.java

@@ -16,13 +16,15 @@
 package org.springframework.data.elasticsearch.repository.query;
 
 import org.springframework.data.elasticsearch.core.ReactiveElasticsearchOperations;
+import org.springframework.data.elasticsearch.core.query.CriteriaQuery;
 import org.springframework.data.elasticsearch.core.query.Query;
 import org.springframework.data.elasticsearch.repository.query.parser.ElasticsearchQueryCreator;
 import org.springframework.data.repository.query.ResultProcessor;
 import org.springframework.data.repository.query.parser.PartTree;
 
 /**
  * @author Christoph Strobl
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
 public class ReactivePartTreeElasticsearchQuery extends AbstractReactiveElasticsearchRepositoryQuery {
@@ -40,7 +42,12 @@ public ReactivePartTreeElasticsearchQuery(ReactiveElasticsearchQueryMethod query
 
 	@Override
 	protected Query createQuery(ElasticsearchParameterAccessor accessor) {
-		return new ElasticsearchQueryCreator(tree, accessor, getMappingContext()).createQuery();
+		CriteriaQuery query = new ElasticsearchQueryCreator(tree, accessor, getMappingContext()).createQuery();
+
+		if (tree.isLimiting()) {
+			query.setMaxResults(tree.getMaxResults());
+		}
+		return query;
 	}
 
 	@Override

src/main/java/org/springframework/data/elasticsearch/repository/query/parser/ElasticsearchQueryCreator.java

@@ -88,7 +88,6 @@ protected CriteriaQuery or(CriteriaQuery base, CriteriaQuery query) {
 	protected CriteriaQuery complete(@Nullable CriteriaQuery query, Sort sort) {
 
 		if (query == null) {
-
 			// this is the case in a findAllByOrderByField method, add empty criteria
 			query = new CriteriaQuery(new Criteria());
 		}