Skip to content

Commit fe8c4f1

Browse files
sothawosothawo
authored
Add @CountQuery annotation.
Original Pull Request #1682 Closes #1156
1 parent 910ca7b commit fe8c4f1

15 files changed

+308
-23
lines changed

‎src/main/java/org/springframework/data/elasticsearch/annotations/CountQuery.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2021 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.elasticsearch.annotations;
17+
18+
import java.lang.annotation.Documented;
19+
import java.lang.annotation.ElementType;
20+
import java.lang.annotation.Retention;
21+
import java.lang.annotation.RetentionPolicy;
22+
import java.lang.annotation.Target;
23+
24+
import org.springframework.core.annotation.AliasFor;
25+
26+
/**
27+
* Alias for a @Query annotation with the count parameter set to true.
28+
*
29+
* @author Peter-Josef Meisch
30+
*/
31+
@Retention(RetentionPolicy.RUNTIME)
32+
@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
33+
@Documented
34+
@Query(count = true)
35+
public @interface CountQuery {
36+
37+
@AliasFor(annotation = Query.class)
38+
String value() default "";
39+
}

‎src/main/java/org/springframework/data/elasticsearch/annotations/Query.java‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,24 +22,32 @@
2222
*
2323
* @author Rizwan Idrees
2424
* @author Mohsin Husen
25+
* @author Peter-Josef Meisch
2526
*/
2627

2728
@Retention(RetentionPolicy.RUNTIME)
28-
@Target(ElementType.METHOD)
29+
@Target({ElementType.METHOD, ElementType.ANNOTATION_TYPE})
2930
@Documented
3031
public @interface Query {
3132

3233
/**
33-
* Elasticsearch query to be used when executing query. May contain placeholders eg. ?0
34-
*
35-
* @return
34+
* @return Elasticsearch query to be used when executing query. May contain placeholders eg. ?0
3635
*/
3736
String value() default "";
3837

3938
/**
4039
* Named Query Named looked up by repository.
4140
*
42-
* @return
41+
* @deprecated since 4.2, not implemented and used anywhere
4342
*/
4443
String name() default "";
44+
45+
/**
46+
* Returns whether the query defined should be executed as count projection.
47+
*
48+
* @return {@literal false} by default.
49+
* @since 4.2
50+
*/
51+
boolean count() default false;
4552
}
53+

‎src/main/java/org/springframework/data/elasticsearch/repository/query/AbstractElasticsearchRepositoryQuery.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
*
2525
* @author Rizwan Idrees
2626
* @author Mohsin Husen
27+
* @author Peter-Josef Meisch
2728
*/
2829

2930
public abstract class AbstractElasticsearchRepositoryQuery implements RepositoryQuery {
@@ -42,4 +43,10 @@ public AbstractElasticsearchRepositoryQuery(ElasticsearchQueryMethod queryMethod
4243
public QueryMethod getQueryMethod() {
4344
return queryMethod;
4445
}
46+
47+
/**
48+
* @return {@literal true} if this is a count query
49+
* @since 4.2
50+
*/
51+
public abstract boolean isCountQuery();
4552
}

‎src/main/java/org/springframework/data/elasticsearch/repository/query/AbstractReactiveElasticsearchRepositoryQuery.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
*/
4444
abstract class AbstractReactiveElasticsearchRepositoryQuery implements RepositoryQuery {
4545

46-
private final ReactiveElasticsearchQueryMethod queryMethod;
46+
protected final ReactiveElasticsearchQueryMethod queryMethod;
4747
private final ReactiveElasticsearchOperations elasticsearchOperations;
4848

4949
AbstractReactiveElasticsearchRepositoryQuery(ReactiveElasticsearchQueryMethod queryMethod,

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,11 @@ public ElasticsearchPartQuery(ElasticsearchQueryMethod method, ElasticsearchOper
5959
this.mappingContext = elasticsearchConverter.getMappingContext();
6060
}
6161

62+
@Override
63+
public boolean isCountQuery() {
64+
return tree.isCountProjection();
65+
}
66+
6267
@Override
6368
public Object execute(Object[] parameters) {
6469
Class<?> clazz = queryMethod.getResultProcessor().getReturnedType().getDomainType();

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchQueryMethod.java‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@
2020
import java.util.Collection;
2121
import java.util.stream.Stream;
2222

23-
import org.springframework.core.annotation.AnnotationUtils;
23+
import org.springframework.core.annotation.AnnotatedElementUtils;
24+
import org.springframework.dao.InvalidDataAccessApiUsageException;
2425
import org.springframework.data.elasticsearch.annotations.Highlight;
2526
import org.springframework.data.elasticsearch.annotations.Query;
2627
import org.springframework.data.elasticsearch.core.SearchHit;
@@ -34,7 +35,9 @@
3435
import org.springframework.data.projection.ProjectionFactory;
3536
import org.springframework.data.repository.core.RepositoryMetadata;
3637
import org.springframework.data.repository.query.QueryMethod;
38+
import org.springframework.data.util.ClassTypeInformation;
3739
import org.springframework.data.util.Lazy;
40+
import org.springframework.data.util.TypeInformation;
3841
import org.springframework.lang.Nullable;
3942
import org.springframework.util.Assert;
4043
import org.springframework.util.ClassUtils;
@@ -53,9 +56,9 @@ public class ElasticsearchQueryMethod extends QueryMethod {
5356

5457
private final MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext;
5558
private @Nullable ElasticsearchEntityMetadata<?> metadata;
56-
private final Method method; // private in base class, but needed here as well
57-
private final Query queryAnnotation;
58-
private final Highlight highlightAnnotation;
59+
protected final Method method; // private in base class, but needed here and in derived classes as well
60+
@Nullable private final Query queryAnnotation;
61+
@Nullable private final Highlight highlightAnnotation;
5962
private final Lazy<HighlightQuery> highlightQueryLazy = Lazy.of(this::createAnnotatedHighlightQuery);
6063

6164
public ElasticsearchQueryMethod(Method method, RepositoryMetadata repositoryMetadata, ProjectionFactory factory,
@@ -67,16 +70,32 @@ public ElasticsearchQueryMethod(Method method, RepositoryMetadata repositoryMeta
6770

6871
this.method = method;
6972
this.mappingContext = mappingContext;
70-
this.queryAnnotation = method.getAnnotation(Query.class);
71-
this.highlightAnnotation = method.getAnnotation(Highlight.class);
73+
this.queryAnnotation = AnnotatedElementUtils.findMergedAnnotation(method, Query.class);
74+
this.highlightAnnotation = AnnotatedElementUtils.findMergedAnnotation(method, Highlight.class);
75+
76+
verifyCountQueryTypes();
77+
}
78+
79+
protected void verifyCountQueryTypes() {
80+
81+
if (hasCountQueryAnnotation()) {
82+
TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
83+
84+
if (returnType.getType() != long.class && !Long.class.isAssignableFrom(returnType.getType())) {
85+
throw new InvalidDataAccessApiUsageException("count query methods must return a Long");
86+
}
87+
}
7288
}
7389

7490
public boolean hasAnnotatedQuery() {
7591
return this.queryAnnotation != null;
7692
}
7793

94+
/**
95+
* @return the query String. Must not be {@literal null} when {@link #hasAnnotatedQuery()} returns true
96+
*/
7897
public String getAnnotatedQuery() {
79-
return (String) AnnotationUtils.getValue(queryAnnotation, "value");
98+
return queryAnnotation.value();
8099
}
81100

82101
/**
@@ -217,4 +236,14 @@ public boolean isNotSearchHitMethod() {
217236
public boolean isNotSearchPageMethod() {
218237
return !isSearchPageMethod();
219238
}
239+
240+
/**
241+
* @return {@literal true} if the method is annotated with
242+
* {@link org.springframework.data.elasticsearch.annotations.CountQuery} or with {@link Query}(count =true)
243+
* @since 4.2
244+
*/
245+
public boolean hasCountQueryAnnotation() {
246+
return queryAnnotation != null && queryAnnotation.count();
247+
}
248+
220249
}

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchStringQuery.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,14 @@ public ElasticsearchStringQuery(ElasticsearchQueryMethod queryMethod, Elasticsea
6969
this.query = query;
7070
}
7171

72+
@Override
73+
public boolean isCountQuery() {
74+
return queryMethod.hasCountQueryAnnotation();
75+
}
76+
7277
@Override
7378
public Object execute(Object[] parameters) {
79+
7480
Class<?> clazz = queryMethod.getResultProcessor().getReturnedType().getDomainType();
7581
ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
7682

@@ -86,7 +92,9 @@ public Object execute(Object[] parameters) {
8692

8793
Object result = null;
8894

89-
if (queryMethod.isPageQuery()) {
95+
if (isCountQuery()) {
96+
result = elasticsearchOperations.count(stringQuery, clazz, index);
97+
} else if (queryMethod.isPageQuery()) {
9098
stringQuery.setPageable(accessor.getPageable());
9199
SearchHits<?> searchHits = elasticsearchOperations.search(stringQuery, clazz, index);
92100
result = SearchHitSupport.searchPageFor(searchHits, stringQuery.getPageable());

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ReactiveElasticsearchQueryMethod.java‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@
1818
import static org.springframework.data.repository.util.ClassUtils.*;
1919

2020
import reactor.core.publisher.Flux;
21+
import reactor.core.publisher.Mono;
2122

2223
import java.lang.reflect.Method;
2324
import java.lang.reflect.ParameterizedType;
25+
import java.util.List;
2426

2527
import org.springframework.dao.InvalidDataAccessApiUsageException;
2628
import org.springframework.data.domain.Page;
@@ -59,7 +61,6 @@ public ReactiveElasticsearchQueryMethod(Method method, RepositoryMetadata metada
5961
if (hasParameterOfType(method, Pageable.class)) {
6062

6163
TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
62-
6364
boolean multiWrapper = ReactiveWrappers.isMultiValueType(returnType.getType());
6465
boolean singleWrapperWithWrappedPageableResult = ReactiveWrappers.isSingleValueType(returnType.getType())
6566
&& (PAGE_TYPE.isAssignableFrom(returnType.getRequiredComponentType())
@@ -87,6 +88,20 @@ public ReactiveElasticsearchQueryMethod(Method method, RepositoryMetadata metada
8788
&& ReactiveWrappers.isMultiValueType(metadata.getReturnType(method).getType()) || super.isCollectionQuery()));
8889
}
8990

91+
@Override
92+
protected void verifyCountQueryTypes() {
93+
if (hasCountQueryAnnotation()) {
94+
TypeInformation<?> returnType = ClassTypeInformation.fromReturnTypeOf(method);
95+
List<TypeInformation<?>> typeArguments = returnType.getTypeArguments();
96+
97+
if (!Mono.class.isAssignableFrom(returnType.getType()) || typeArguments.size() != 1
98+
|| (typeArguments.get(0).getType() != long.class
99+
&& !Long.class.isAssignableFrom(typeArguments.get(0).getType()))) {
100+
throw new InvalidDataAccessApiUsageException("count query methods must return a Mono<Long>");
101+
}
102+
}
103+
}
104+
90105
@Override
91106
protected ElasticsearchParameters createParameters(Method method) {
92107
return new ElasticsearchParameters(method);

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ReactiveElasticsearchStringQuery.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ private String getParameterWithIndex(ElasticsearchParameterAccessor accessor, in
7575

7676
@Override
7777
boolean isCountQuery() {
78-
return false;
78+
return queryMethod.hasCountQueryAnnotation();
7979
}
8080

8181
@Override

‎src/test/java/org/springframework/data/elasticsearch/core/mapping/FieldNamingStrategyIntegrationTemplateTest.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
11
/*
2-
* (c) Copyright 2021 sothawo
2+
* Copyright 2021 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
315
*/
416
package org.springframework.data.elasticsearch.core.mapping;
517

@@ -10,7 +22,7 @@
1022
import org.springframework.test.context.ContextConfiguration;
1123

1224
/**
13-
* @author P.J. Meisch ([email protected])
25+
* @author Peter-Josef Meisch
1426
*/
1527
@ContextConfiguration(classes = { FieldNamingStrategyIntegrationTemplateTest.Config.class })
1628
public class FieldNamingStrategyIntegrationTemplateTest extends FieldNamingStrategyIntegrationTest {

0 commit comments

Comments
 (0)