Skip to content

Commit efd592e

Browse files
sothawosothawo
committed
DATAES-631 - Added QueryMapper classes.
1 parent b778d5b commit efd592e

8 files changed

+447
-73
lines changed

‎src/main/java/org/springframework/data/elasticsearch/core/AbstractQueryMapper.java

+120
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*
2+
* Copyright 2019 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.core;
17+
18+
import org.elasticsearch.index.query.QueryBuilder;
19+
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
20+
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
21+
import org.springframework.data.elasticsearch.core.query.CriteriaQuery;
22+
import org.springframework.data.elasticsearch.core.query.Query;
23+
import org.springframework.data.elasticsearch.core.query.SearchQuery;
24+
import org.springframework.lang.Nullable;
25+
26+
/**
27+
* @author Peter-Josef Meisch
28+
* @since 4.0
29+
*/
30+
public abstract class AbstractQueryMapper<R> {
31+
32+
protected ElasticsearchConverter elasticsearchConverter;
33+
34+
public AbstractQueryMapper(ElasticsearchConverter elasticsearchConverter) {
35+
this.elasticsearchConverter = elasticsearchConverter;
36+
}
37+
38+
public abstract <T> R mapQuery(Query query, Class<T> clazz);
39+
40+
protected Pair<QueryBuilder, QueryBuilder> getQueryAndFilter(Query query) {
41+
QueryBuilder elasticsearchQuery;
42+
QueryBuilder elasticsearchFilter;
43+
44+
if (query instanceof SearchQuery) {
45+
SearchQuery searchQuery = (SearchQuery) query;
46+
elasticsearchQuery = searchQuery.getQuery();
47+
elasticsearchFilter = searchQuery.getFilter();
48+
} else if (query instanceof CriteriaQuery) {
49+
CriteriaQuery criteriaQuery = (CriteriaQuery) query;
50+
elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());
51+
elasticsearchFilter = new CriteriaFilterProcessor().createFilterFromCriteria(criteriaQuery.getCriteria());
52+
} else {
53+
throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
54+
}
55+
56+
return new Pair<>(elasticsearchQuery, elasticsearchFilter);
57+
}
58+
59+
@Nullable
60+
protected String[] retrieveIndexNameFromPersistentEntity(@Nullable Class clazz) {
61+
return (clazz != null) ? new String[] { getPersistentEntityFor(clazz).getIndexName() } : null;
62+
}
63+
64+
@Nullable
65+
protected String[] retrieveTypeFromPersistentEntity(@Nullable Class clazz) {
66+
return (clazz != null) ? new String[] { getPersistentEntityFor(clazz).getIndexType() } : null;
67+
}
68+
69+
protected ElasticsearchPersistentEntity getPersistentEntityFor(Class clazz) {
70+
return elasticsearchConverter.getMappingContext().getRequiredPersistentEntity(clazz);
71+
}
72+
73+
@Nullable
74+
protected ElasticsearchPersistentEntity<?> getPersistentEntity(@Nullable Class<?> clazz) {
75+
return clazz != null ? elasticsearchConverter.getMappingContext().getPersistentEntity(clazz) : null;
76+
}
77+
78+
protected void setPersistentEntityIndexAndType(Query query, @Nullable Class clazz) {
79+
if (query.getIndices().isEmpty()) {
80+
String[] indices = retrieveIndexNameFromPersistentEntity(clazz);
81+
82+
if (indices != null) {
83+
query.addIndices(indices);
84+
}
85+
}
86+
if (query.getTypes().isEmpty()) {
87+
String[] types = retrieveTypeFromPersistentEntity(clazz);
88+
89+
if (types != null) {
90+
query.addTypes(types);
91+
}
92+
}
93+
}
94+
95+
/**
96+
* Pair with potentially null elements.
97+
*
98+
* @param <F>
99+
* @param <S>
100+
*/
101+
static class Pair<F, S> {
102+
private final F first;
103+
private final S second;
104+
105+
public Pair(@Nullable F first, @Nullable S second) {
106+
this.first = first;
107+
this.second = second;
108+
}
109+
110+
@Nullable
111+
public F getFirst() {
112+
return first;
113+
}
114+
115+
@Nullable
116+
public S getSecond() {
117+
return second;
118+
}
119+
}
120+
}

‎src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,7 @@ default List<List<?>> queryForList(List<SearchQuery> queries, List<Class<?>> cla
330330
* @param clazz
331331
* @return
332332
*/
333-
<T> long count(Query query, Class<T> clazz);
333+
<T> long count(Query query, @Nullable Class<T> clazz);
334334

335335
/**
336336
* return number of elements found by given query

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

+4-32
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,8 @@ public class ElasticsearchRestTemplate extends AbstractElasticsearchTemplate
151151

152152
private static final Logger logger = LoggerFactory.getLogger(ElasticsearchRestTemplate.class);
153153

154+
private ElasticsearchRestTemplateQueryMapper queryMapper;
155+
154156
private RestHighLevelClient client;
155157
private String searchTimeout;
156158

@@ -169,6 +171,7 @@ private void initialize(RestHighLevelClient client, ElasticsearchConverter elast
169171

170172
this.client = client;
171173
this.elasticsearchConverter = elasticsearchConverter;
174+
this.queryMapper = new ElasticsearchRestTemplateQueryMapper(elasticsearchConverter);
172175
}
173176

174177
@Override
@@ -514,7 +517,7 @@ private <T> CloseableIterator<T> doStream(long scrollTimeInMillis, ScrolledPage<
514517

515518
@Override
516519
public <T> long count(Query query, Class<T> clazz) {
517-
SearchRequest searchRequest = mapQuery(query, clazz);
520+
SearchRequest searchRequest = queryMapper.mapQuery(query, clazz);
518521
searchRequest.source().size(0);
519522

520523
try {
@@ -524,37 +527,6 @@ public <T> long count(Query query, Class<T> clazz) {
524527
}
525528
}
526529

527-
private <T> SearchRequest mapQuery(Query query, Class<T> clazz) {
528-
QueryBuilder elasticsearchQuery;
529-
QueryBuilder elasticsearchFilter;
530-
531-
if (query instanceof SearchQuery) {
532-
SearchQuery searchQuery = (SearchQuery) query;
533-
elasticsearchQuery = searchQuery.getQuery();
534-
elasticsearchFilter = searchQuery.getFilter();
535-
} else if (query instanceof CriteriaQuery) {
536-
CriteriaQuery criteriaQuery = (CriteriaQuery) query;
537-
elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());
538-
elasticsearchFilter = new CriteriaFilterProcessor().createFilterFromCriteria(criteriaQuery.getCriteria());
539-
} else {
540-
throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
541-
}
542-
543-
SearchRequest searchRequest = prepareSearch(query, clazz);
544-
545-
if (elasticsearchQuery != null) {
546-
searchRequest.source().query(elasticsearchQuery);
547-
} else {
548-
searchRequest.source().query(QueryBuilders.matchAllQuery());
549-
}
550-
551-
if (elasticsearchFilter != null) {
552-
searchRequest.source().postFilter(elasticsearchFilter);
553-
}
554-
555-
return searchRequest;
556-
}
557-
558530
@Override
559531
public <T> List<T> multiGet(SearchQuery searchQuery, Class<T> clazz) {
560532
return elasticsearchConverter.mapDocuments(DocumentAdapters.from(getMultiResponse(searchQuery, clazz)), clazz);

‎src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplateQueryMapper.java

+160
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/*
2+
* Copyright 2019 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.core;
17+
18+
import java.util.Optional;
19+
20+
import org.elasticsearch.action.search.SearchRequest;
21+
import org.elasticsearch.index.query.QueryBuilder;
22+
import org.elasticsearch.index.query.QueryBuilders;
23+
import org.elasticsearch.search.builder.SearchSourceBuilder;
24+
import org.elasticsearch.search.sort.FieldSortBuilder;
25+
import org.elasticsearch.search.sort.ScoreSortBuilder;
26+
import org.elasticsearch.search.sort.SortBuilders;
27+
import org.elasticsearch.search.sort.SortOrder;
28+
import org.springframework.data.domain.Sort;
29+
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
30+
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
31+
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
32+
import org.springframework.data.elasticsearch.core.query.Query;
33+
import org.springframework.data.elasticsearch.core.query.SourceFilter;
34+
import org.springframework.lang.Nullable;
35+
import org.springframework.util.Assert;
36+
37+
/**
38+
* @author Peter-Josef Meisch
39+
* @since 4.0
40+
*/
41+
public class ElasticsearchRestTemplateQueryMapper extends AbstractQueryMapper<SearchRequest> {
42+
43+
public ElasticsearchRestTemplateQueryMapper(ElasticsearchConverter elasticsearchConverter) {
44+
super(elasticsearchConverter);
45+
}
46+
47+
@Override
48+
public <T> SearchRequest mapQuery(Query query, Class<T> clazz) {
49+
SearchRequest searchRequest = prepareSearch(query, clazz);
50+
51+
Pair<QueryBuilder, QueryBuilder> queryAndFilter = getQueryAndFilter(query);
52+
QueryBuilder elasticsearchQuery = queryAndFilter.getFirst();
53+
QueryBuilder elasticsearchFilter = queryAndFilter.getSecond();
54+
55+
if (elasticsearchQuery != null) {
56+
searchRequest.source().query(elasticsearchQuery);
57+
} else {
58+
searchRequest.source().query(QueryBuilders.matchAllQuery());
59+
}
60+
61+
if (elasticsearchFilter != null) {
62+
searchRequest.source().postFilter(elasticsearchFilter);
63+
}
64+
65+
return searchRequest;
66+
}
67+
68+
protected <T> SearchRequest prepareSearch(Query query, @Nullable Class<T> clazz) {
69+
setPersistentEntityIndexAndType(query, clazz);
70+
return prepareSearch(query, Optional.empty(), clazz);
71+
}
72+
73+
private SearchRequest prepareSearch(Query query, Optional<QueryBuilder> builder, @Nullable Class<?> clazz) {
74+
Assert.notNull(query.getIndices(), "No index defined for Query");
75+
Assert.notEmpty(query.getIndices(), "No index defined for Query");
76+
Assert.notNull(query.getTypes(), "No type defined for Query");
77+
78+
int startRecord = 0;
79+
SearchRequest request = new SearchRequest(query.getIndices().toArray(new String[0]));
80+
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
81+
request.types(query.getTypes().toArray(new String[0]));
82+
sourceBuilder.version(true);
83+
sourceBuilder.trackScores(query.getTrackScores());
84+
85+
if (builder.isPresent()) {
86+
sourceBuilder.query(builder.get());
87+
}
88+
89+
if (query.getSourceFilter() != null) {
90+
SourceFilter sourceFilter = query.getSourceFilter();
91+
sourceBuilder.fetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
92+
}
93+
94+
if (query.getPageable().isPaged()) {
95+
startRecord = query.getPageable().getPageNumber() * query.getPageable().getPageSize();
96+
sourceBuilder.size(query.getPageable().getPageSize());
97+
}
98+
sourceBuilder.from(startRecord);
99+
100+
if (!query.getFields().isEmpty()) {
101+
sourceBuilder.fetchSource(query.getFields().toArray(new String[0]), null);
102+
}
103+
104+
if (query.getIndicesOptions() != null) {
105+
request.indicesOptions(query.getIndicesOptions());
106+
}
107+
108+
if (query.getSort() != null) {
109+
prepareSort(query, sourceBuilder, getPersistentEntity(clazz));
110+
}
111+
112+
if (query.getMinScore() > 0) {
113+
sourceBuilder.minScore(query.getMinScore());
114+
}
115+
116+
if (query.getPreference() != null) {
117+
request.preference(query.getPreference());
118+
}
119+
120+
if (query.getSearchType() != null) {
121+
request.searchType(query.getSearchType());
122+
}
123+
124+
request.source(sourceBuilder);
125+
return request;
126+
}
127+
128+
private void prepareSort(Query query, SearchSourceBuilder sourceBuilder,
129+
@Nullable ElasticsearchPersistentEntity<?> entity) {
130+
131+
for (Sort.Order order : query.getSort()) {
132+
SortOrder sortOrder = order.getDirection().isDescending() ? SortOrder.DESC : SortOrder.ASC;
133+
134+
if (ScoreSortBuilder.NAME.equals(order.getProperty())) {
135+
ScoreSortBuilder sort = SortBuilders //
136+
.scoreSort() //
137+
.order(sortOrder);
138+
139+
sourceBuilder.sort(sort);
140+
} else {
141+
ElasticsearchPersistentProperty property = (entity != null) //
142+
? entity.getPersistentProperty(order.getProperty()) //
143+
: null;
144+
String fieldName = property != null ? property.getFieldName() : order.getProperty();
145+
146+
FieldSortBuilder sort = SortBuilders //
147+
.fieldSort(fieldName) //
148+
.order(sortOrder);
149+
150+
if (order.getNullHandling() == Sort.NullHandling.NULLS_FIRST) {
151+
sort.missing("_first");
152+
} else if (order.getNullHandling() == Sort.NullHandling.NULLS_LAST) {
153+
sort.missing("_last");
154+
}
155+
156+
sourceBuilder.sort(sort);
157+
}
158+
}
159+
}
160+
}

0 commit comments

Comments
 (0)