Add @CountQuery annotation.

Original Pull Request #1682 
Closes #1156

Author: sothawo

Diff for: src/main/java/org/springframework/data/elasticsearch/annotations/CountQuery.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.elasticsearch.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.springframework.core.annotation.AliasFor;
+
+/**
+ * Alias for a @Query annotation with the count parameter set to true.
+ *
+ * @author Peter-Josef Meisch
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+@Documented
+@Query(count = true)
+public @interface CountQuery {
+
+	@AliasFor(annotation = Query.class)
+	String value() default "";
+}

Diff for: src/main/java/org/springframework/data/elasticsearch/annotations/Query.java

@@ -22,24 +22,32 @@
  *
  * @author Rizwan Idrees
  * @author Mohsin Husen
+ * @author Peter-Josef Meisch
  */
 
 @Retention(RetentionPolicy.RUNTIME)
-@Target(ElementType.METHOD)
+@Target({ElementType.METHOD, ElementType.ANNOTATION_TYPE})
 @Documented
 public @interface Query {
 
 	/**
-	 * Elasticsearch query to be used when executing query. May contain placeholders eg. ?0
-	 *
-	 * @return
+	 * @return Elasticsearch query to be used when executing query. May contain placeholders eg. ?0
 	 */
 	String value() default "";
 
 	/**
 	 * Named Query Named looked up by repository.
 	 *
-	 * @return
+	 * @deprecated since 4.2, not implemented and used anywhere
 	 */
 	String name() default "";
+
+	/**
+	 * Returns whether the query defined should be executed as count projection.
+	 *
+	 * @return {@literal false} by default.
+	 * @since 4.2
+	 */
+	boolean count() default false;
 }
+

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/AbstractElasticsearchRepositoryQuery.java

@@ -24,6 +24,7 @@
  *
  * @author Rizwan Idrees
  * @author Mohsin Husen
+ * @author Peter-Josef Meisch
  */
 
 public abstract class AbstractElasticsearchRepositoryQuery implements RepositoryQuery {
@@ -42,4 +43,10 @@ public AbstractElasticsearchRepositoryQuery(ElasticsearchQueryMethod queryMethod
 	public QueryMethod getQueryMethod() {
 		return queryMethod;
 	}
+
+	/**
+	 * @return {@literal true} if this is a count query
+	 * @since 4.2
+	 */
+	public abstract boolean isCountQuery();
 }

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/AbstractReactiveElasticsearchRepositoryQuery.java

@@ -43,7 +43,7 @@
  */
 abstract class AbstractReactiveElasticsearchRepositoryQuery implements RepositoryQuery {
 
-	private final ReactiveElasticsearchQueryMethod queryMethod;
+	protected final ReactiveElasticsearchQueryMethod queryMethod;
 	private final ReactiveElasticsearchOperations elasticsearchOperations;
 
 	AbstractReactiveElasticsearchRepositoryQuery(ReactiveElasticsearchQueryMethod queryMethod,

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

@@ -59,6 +59,11 @@ public ElasticsearchPartQuery(ElasticsearchQueryMethod method, ElasticsearchOper
 		this.mappingContext = elasticsearchConverter.getMappingContext();
 	}
 
+	@Override
+	public boolean isCountQuery() {
+		return tree.isCountProjection();
+	}
+
 	@Override
 	public Object execute(Object[] parameters) {
 		Class<?> clazz = queryMethod.getResultProcessor().getReturnedType().getDomainType();

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchQueryMethod.java

@@ -20,7 +20,8 @@
 import java.util.Collection;
 import java.util.stream.Stream;
 
-import org.springframework.core.annotation.AnnotationUtils;
+import org.springframework.core.annotation.AnnotatedElementUtils;
+import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.elasticsearch.annotations.Highlight;
 import org.springframework.data.elasticsearch.annotations.Query;
 import org.springframework.data.elasticsearch.core.SearchHit;
@@ -34,7 +35,9 @@
 import org.springframework.data.projection.ProjectionFactory;
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.query.QueryMethod;
+import org.springframework.data.util.ClassTypeInformation;
 import org.springframework.data.util.Lazy;
+import org.springframework.data.util.TypeInformation;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
@@ -53,9 +56,9 @@ public class ElasticsearchQueryMethod extends QueryMethod {
 
 	private final MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext;
 	private @Nullable ElasticsearchEntityMetadata<?> metadata;
-	private final Method method; // private in base class, but needed here as well
-	private final Query queryAnnotation;
-	private final Highlight highlightAnnotation;
+	protected final Method method; // private in base class, but needed here and in derived classes as well
+	@Nullable private final Query queryAnnotation;
+	@Nullable private final Highlight highlightAnnotation;
 	private final Lazy<HighlightQuery> highlightQueryLazy = Lazy.of(this::createAnnotatedHighlightQuery);
 
 	public ElasticsearchQueryMethod(Method method, RepositoryMetadata repositoryMetadata, ProjectionFactory factory,
@@ -67,16 +70,32 @@ public ElasticsearchQueryMethod(Method method, RepositoryMetadata repositoryMeta
 
 		this.method = method;
 		this.mappingContext = mappingContext;
-		this.queryAnnotation = method.getAnnotation(Query.class);
-		this.highlightAnnotation = method.getAnnotation(Highlight.class);
+		this.queryAnnotation = AnnotatedElementUtils.findMergedAnnotation(method, Query.class);
+		this.highlightAnnotation = AnnotatedElementUtils.findMergedAnnotation(method, Highlight.class);
+
+		verifyCountQueryTypes();
+	}
+
+	protected void verifyCountQueryTypes() {
+
+		if (hasCountQueryAnnotation()) {
+			TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
+
+			if (returnType.getType() != long.class && !Long.class.isAssignableFrom(returnType.getType())) {
+				throw new InvalidDataAccessApiUsageException("count query methods must return a Long");
+			}
+		}
 	}
 
 	public boolean hasAnnotatedQuery() {
 		return this.queryAnnotation != null;
 	}
 
+	/**
+	 * @return the query String. Must not be {@literal null} when {@link #hasAnnotatedQuery()} returns true
+	 */
 	public String getAnnotatedQuery() {
-		return (String) AnnotationUtils.getValue(queryAnnotation, "value");
+		return queryAnnotation.value();
 	}
 
 	/**
@@ -217,4 +236,14 @@ public boolean isNotSearchHitMethod() {
 	public boolean isNotSearchPageMethod() {
 		return !isSearchPageMethod();
 	}
+
+	/**
+	 * @return {@literal true} if the method is annotated with
+	 *         {@link org.springframework.data.elasticsearch.annotations.CountQuery} or with {@link Query}(count =true)
+	 * @since 4.2
+	 */
+	public boolean hasCountQueryAnnotation() {
+		return queryAnnotation != null && queryAnnotation.count();
+	}
+
 }

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchStringQuery.java

@@ -69,8 +69,14 @@ public ElasticsearchStringQuery(ElasticsearchQueryMethod queryMethod, Elasticsea
 		this.query = query;
 	}
 
+	@Override
+	public boolean isCountQuery() {
+		return queryMethod.hasCountQueryAnnotation();
+	}
+
 	@Override
 	public Object execute(Object[] parameters) {
+
 		Class<?> clazz = queryMethod.getResultProcessor().getReturnedType().getDomainType();
 		ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
 
@@ -86,7 +92,9 @@ public Object execute(Object[] parameters) {
 
 		Object result = null;
 
-		if (queryMethod.isPageQuery()) {
+		if (isCountQuery()) {
+			result = elasticsearchOperations.count(stringQuery, clazz, index);
+		} else if (queryMethod.isPageQuery()) {
 			stringQuery.setPageable(accessor.getPageable());
 			SearchHits<?> searchHits = elasticsearchOperations.search(stringQuery, clazz, index);
 			result = SearchHitSupport.searchPageFor(searchHits, stringQuery.getPageable());

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/ReactiveElasticsearchQueryMethod.java

@@ -18,9 +18,11 @@
 import static org.springframework.data.repository.util.ClassUtils.*;
 
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import java.lang.reflect.Method;
 import java.lang.reflect.ParameterizedType;
+import java.util.List;
 
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.domain.Page;
@@ -59,7 +61,6 @@ public ReactiveElasticsearchQueryMethod(Method method, RepositoryMetadata metada
 		if (hasParameterOfType(method, Pageable.class)) {
 
 			TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
-
 			boolean multiWrapper = ReactiveWrappers.isMultiValueType(returnType.getType());
 			boolean singleWrapperWithWrappedPageableResult = ReactiveWrappers.isSingleValueType(returnType.getType())
 					&& (PAGE_TYPE.isAssignableFrom(returnType.getRequiredComponentType())
@@ -87,6 +88,20 @@ public ReactiveElasticsearchQueryMethod(Method method, RepositoryMetadata metada
 				&& ReactiveWrappers.isMultiValueType(metadata.getReturnType(method).getType()) || super.isCollectionQuery()));
 	}
 
+	@Override
+	protected void verifyCountQueryTypes() {
+		if (hasCountQueryAnnotation()) {
+			TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
+			List<TypeInformation<?>> typeArguments = returnType.getTypeArguments();
+
+			if (!Mono.class.isAssignableFrom(returnType.getType()) || typeArguments.size() != 1
+					|| (typeArguments.get(0).getType() != long.class
+							&& !Long.class.isAssignableFrom(typeArguments.get(0).getType()))) {
+				throw new InvalidDataAccessApiUsageException("count query methods must return a Mono<Long>");
+			}
+		}
+	}
+
 	@Override
 	protected ElasticsearchParameters createParameters(Method method) {
 		return new ElasticsearchParameters(method);

Diff for: src/main/java/org/springframework/data/elasticsearch/repository/query/ReactiveElasticsearchStringQuery.java

@@ -75,7 +75,7 @@ private String getParameterWithIndex(ElasticsearchParameterAccessor accessor, in
 
 	@Override
 	boolean isCountQuery() {
-		return false;
+		return queryMethod.hasCountQueryAnnotation();
 	}
 
 	@Override

Diff for: src/test/java/org/springframework/data/elasticsearch/core/mapping/FieldNamingStrategyIntegrationTemplateTest.java

@@ -1,5 +1,17 @@
 /*
- * (c) Copyright 2021 sothawo
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 package org.springframework.data.elasticsearch.core.mapping;
 
@@ -10,7 +22,7 @@
 import org.springframework.test.context.ContextConfiguration;
 
 /**
- * @author P.J. Meisch ([email protected])
+ * @author Peter-Josef Meisch
  */
 @ContextConfiguration(classes = { FieldNamingStrategyIntegrationTemplateTest.Config.class })
 public class FieldNamingStrategyIntegrationTemplateTest extends FieldNamingStrategyIntegrationTest {

Diff for: src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java

@@ -35,6 +35,7 @@
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.annotation.Id;
@@ -44,6 +45,7 @@
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Order;
+import org.springframework.data.elasticsearch.annotations.CountQuery;
 import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.annotations.Field;
 import org.springframework.data.elasticsearch.annotations.Highlight;
@@ -911,6 +913,31 @@ public void shouldCountCustomMethod() {
 		assertThat(count).isEqualTo(1L);
 	}
 
+	@Test // #1156
+	@DisplayName("should count with query by type")
+	void shouldCountWithQueryByType() {
+
+		String documentId = nextIdAsString();
+		SampleEntity sampleEntity = new SampleEntity();
+		sampleEntity.setId(documentId);
+		sampleEntity.setType("test");
+		sampleEntity.setMessage("some message");
+
+		repository.save(sampleEntity);
+
+		documentId = nextIdAsString();
+		SampleEntity sampleEntity2 = new SampleEntity();
+		sampleEntity2.setId(documentId);
+		sampleEntity2.setType("test2");
+		sampleEntity2.setMessage("some message");
+
+		repository.save(sampleEntity2);
+
+		long count = repository.countWithQueryByType("test");
+
+		assertThat(count).isEqualTo(1L);
+	}
+
 	@Test // DATAES-106
 	public void shouldCountCustomMethodForNot() {
 
@@ -1746,6 +1773,9 @@ public interface SampleCustomMethodRepository extends ElasticsearchRepository<Sa
 		SearchHits<SampleEntity> searchBy(Sort sort);
 
 		SearchPage<SampleEntity> searchByMessage(String message, Pageable pageable);
+
+		@CountQuery("{\"bool\" : {\"must\" : {\"term\" : {\"type\" : \"?0\"}}}}")
+		long countWithQueryByType(String type);
 	}
 
 	/**