1 | /* |
2 | * Copyright 2013 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 | * http://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 | |
17 | package org.springframework.data.elasticsearch.core; |
18 | |
19 | import org.codehaus.jackson.map.DeserializationConfig; |
20 | import org.codehaus.jackson.map.ObjectMapper; |
21 | import org.elasticsearch.action.bulk.BulkItemResponse; |
22 | import org.elasticsearch.action.bulk.BulkRequestBuilder; |
23 | import org.elasticsearch.action.bulk.BulkResponse; |
24 | import org.elasticsearch.action.count.CountRequestBuilder; |
25 | import org.elasticsearch.action.get.GetResponse; |
26 | import org.elasticsearch.action.index.IndexRequestBuilder; |
27 | import org.elasticsearch.action.search.SearchRequestBuilder; |
28 | import org.elasticsearch.action.search.SearchResponse; |
29 | import org.elasticsearch.client.Client; |
30 | import org.elasticsearch.client.Requests; |
31 | import org.elasticsearch.common.collect.MapBuilder; |
32 | import org.elasticsearch.common.unit.TimeValue; |
33 | import org.elasticsearch.index.query.FilterBuilder; |
34 | import org.elasticsearch.index.query.QueryBuilder; |
35 | import org.elasticsearch.search.SearchHit; |
36 | import org.elasticsearch.search.sort.SortBuilder; |
37 | import org.elasticsearch.search.sort.SortOrder; |
38 | import org.springframework.data.domain.Page; |
39 | import org.springframework.data.domain.PageImpl; |
40 | import org.springframework.data.domain.Pageable; |
41 | import org.springframework.data.domain.Sort; |
42 | import org.springframework.data.elasticsearch.ElasticsearchException; |
43 | import org.springframework.data.elasticsearch.annotations.Document; |
44 | import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter; |
45 | import org.springframework.data.elasticsearch.core.convert.MappingElasticsearchConverter; |
46 | import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity; |
47 | import org.springframework.data.elasticsearch.core.mapping.SimpleElasticsearchMappingContext; |
48 | import org.springframework.data.elasticsearch.core.query.*; |
49 | import org.springframework.util.Assert; |
50 | |
51 | import java.io.IOException; |
52 | import java.util.ArrayList; |
53 | import java.util.HashMap; |
54 | import java.util.List; |
55 | import java.util.Map; |
56 | |
57 | import static org.apache.commons.lang.StringUtils.isBlank; |
58 | import static org.elasticsearch.action.search.SearchType.DFS_QUERY_THEN_FETCH; |
59 | import static org.elasticsearch.action.search.SearchType.SCAN; |
60 | import static org.elasticsearch.client.Requests.indicesExistsRequest; |
61 | import static org.elasticsearch.client.Requests.refreshRequest; |
62 | import static org.elasticsearch.index.VersionType.EXTERNAL; |
63 | |
64 | /** |
65 | * ElasticsearchTemplate |
66 | * |
67 | * @author Rizwan Idrees |
68 | * @author Mohsin Husen |
69 | */ |
70 | |
71 | public class ElasticsearchTemplate implements ElasticsearchOperations { |
72 | |
73 | private Client client; |
74 | private ElasticsearchConverter elasticsearchConverter; |
75 | private ObjectMapper objectMapper = new ObjectMapper(); |
76 | |
77 | { |
78 | objectMapper.configure(DeserializationConfig.Feature.FAIL_ON_UNKNOWN_PROPERTIES, false); |
79 | } |
80 | |
81 | public ElasticsearchTemplate(Client client) { |
82 | this(client, null); |
83 | } |
84 | |
85 | public ElasticsearchTemplate(Client client, ElasticsearchConverter elasticsearchConverter) { |
86 | this.client = client; |
87 | this.elasticsearchConverter = (elasticsearchConverter == null)? new MappingElasticsearchConverter(new SimpleElasticsearchMappingContext()) : elasticsearchConverter ; |
88 | } |
89 | |
90 | |
91 | @Override |
92 | public <T> boolean createIndex(Class<T> clazz) { |
93 | ElasticsearchPersistentEntity<T> persistentEntity = getPersistentEntityFor(clazz); |
94 | return createIndexIfNotCreated(persistentEntity.getIndexName()); |
95 | } |
96 | |
97 | @Override |
98 | public ElasticsearchConverter getElasticsearchConverter() { |
99 | return elasticsearchConverter; |
100 | } |
101 | |
102 | @Override |
103 | public <T> T queryForObject(GetQuery query, Class<T> clazz) { |
104 | ElasticsearchPersistentEntity<T> persistentEntity = getPersistentEntityFor(clazz); |
105 | GetResponse response = client.prepareGet(persistentEntity.getIndexName(), persistentEntity.getIndexType(), query.getId()) |
106 | .execute().actionGet(); |
107 | return mapResult(response.getSourceAsString(), clazz); |
108 | } |
109 | |
110 | @Override |
111 | public <T> T queryForObject(CriteriaQuery query, Class<T> clazz) { |
112 | Page<T> page = queryForPage(query,clazz); |
113 | Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found "+ page.getTotalElements() +" results"); |
114 | return page.getTotalElements() > 0? page.getContent().get(0) : null; |
115 | } |
116 | |
117 | @Override |
118 | public <T> T queryForObject(StringQuery query, Class<T> clazz) { |
119 | Page<T> page = queryForPage(query,clazz); |
120 | Assert.isTrue(page.getTotalElements() < 2, "Expected 1 but found "+ page.getTotalElements() +" results"); |
121 | return page.getTotalElements() > 0? page.getContent().get(0) : null; |
122 | } |
123 | |
124 | @Override |
125 | public <T> Page<T> queryForPage(SearchQuery query, Class<T> clazz) { |
126 | SearchResponse response = doSearch(prepareSearch(query,clazz), query.getElasticsearchQuery(), query.getElasticsearchFilter(),query.getElasticsearchSort()); |
127 | return mapResults(response, clazz, query.getPageable()); |
128 | } |
129 | |
130 | |
131 | public <T> Page<T> queryForPage(SearchQuery query, ResultsMapper<T> resultsMapper) { |
132 | SearchResponse response = doSearch(prepareSearch(query), query.getElasticsearchQuery(), query.getElasticsearchFilter(),query.getElasticsearchSort()); |
133 | return resultsMapper.mapResults(response); |
134 | } |
135 | |
136 | @Override |
137 | public <T> List<String> queryForIds(SearchQuery query) { |
138 | SearchRequestBuilder request = prepareSearch(query).setQuery(query.getElasticsearchQuery()) |
139 | .setNoFields(); |
140 | if(query.getElasticsearchFilter() != null){ |
141 | request.setFilter(query.getElasticsearchFilter()); |
142 | } |
143 | SearchResponse response = request.execute().actionGet(); |
144 | return extractIds(response); |
145 | } |
146 | |
147 | private SearchResponse doSearch(SearchRequestBuilder searchRequest, QueryBuilder query, FilterBuilder filter, SortBuilder sortBuilder){ |
148 | if(filter != null){ |
149 | searchRequest.setFilter(filter); |
150 | } |
151 | |
152 | if(sortBuilder != null){ |
153 | searchRequest.addSort(sortBuilder); |
154 | } |
155 | |
156 | return searchRequest.setQuery(query).execute().actionGet(); |
157 | } |
158 | |
159 | @Override |
160 | public <T> Page<T> queryForPage(CriteriaQuery query, Class<T> clazz) { |
161 | QueryBuilder elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(query.getCriteria()); |
162 | SearchResponse response = prepareSearch(query,clazz) |
163 | .setQuery(elasticsearchQuery) |
164 | .execute().actionGet(); |
165 | return mapResults(response, clazz, query.getPageable()); |
166 | } |
167 | |
168 | @Override |
169 | public <T> Page<T> queryForPage(StringQuery query, Class<T> clazz) { |
170 | SearchResponse response = prepareSearch(query,clazz) |
171 | .setQuery(query.getSource()) |
172 | .execute().actionGet(); |
173 | return mapResults(response, clazz, query.getPageable()); |
174 | } |
175 | |
176 | @Override |
177 | public <T> long count(SearchQuery query, Class<T> clazz) { |
178 | ElasticsearchPersistentEntity<T> persistentEntity = getPersistentEntityFor(clazz); |
179 | CountRequestBuilder countRequestBuilder = client.prepareCount(persistentEntity.getIndexName()) |
180 | .setTypes(persistentEntity.getIndexType()); |
181 | if(query.getElasticsearchQuery() != null){ |
182 | countRequestBuilder.setQuery(query.getElasticsearchQuery()); |
183 | } |
184 | return countRequestBuilder.execute().actionGet().count(); |
185 | } |
186 | |
187 | @Override |
188 | public String index(IndexQuery query) { |
189 | return prepareIndex(query) |
190 | .execute() |
191 | .actionGet().getId(); |
192 | } |
193 | |
194 | @Override |
195 | public void bulkIndex(List<IndexQuery> queries) { |
196 | BulkRequestBuilder bulkRequest = client.prepareBulk(); |
197 | for(IndexQuery query : queries){ |
198 | bulkRequest.add(prepareIndex(query)); |
199 | } |
200 | BulkResponse bulkResponse = bulkRequest.execute().actionGet(); |
201 | if (bulkResponse.hasFailures()) { |
202 | Map<String, String> failedDocuments = new HashMap<String, String>(); |
203 | for (BulkItemResponse item : bulkResponse.items()) { |
204 | if (item.failed()) |
205 | failedDocuments.put(item.getId(), item.failureMessage()); |
206 | } |
207 | throw new ElasticsearchException("Bulk indexing has failures. Use ElasticsearchException.getFailedDocuments() for detailed messages [" + failedDocuments+"]", failedDocuments); |
208 | } |
209 | } |
210 | |
211 | @Override |
212 | public String delete(String indexName, String type, String id) { |
213 | return client.prepareDelete(indexName, type, id) |
214 | .execute().actionGet().getId(); |
215 | } |
216 | |
217 | @Override |
218 | public <T> String delete(Class<T> clazz, String id) { |
219 | ElasticsearchPersistentEntity persistentEntity = getPersistentEntityFor(clazz); |
220 | return delete(persistentEntity.getIndexName(), persistentEntity.getIndexType(), id); |
221 | } |
222 | |
223 | @Override |
224 | public <T> void delete(DeleteQuery query, Class<T> clazz) { |
225 | ElasticsearchPersistentEntity persistentEntity = getPersistentEntityFor(clazz); |
226 | client.prepareDeleteByQuery(persistentEntity.getIndexName()) |
227 | .setTypes(persistentEntity.getIndexType()) |
228 | .setQuery(query.getElasticsearchQuery()) |
229 | .execute().actionGet(); |
230 | } |
231 | |
232 | private boolean createIndexIfNotCreated(String indexName) { |
233 | return indexExists(indexName) || createIndex(indexName); |
234 | } |
235 | |
236 | private boolean indexExists(String indexName) { |
237 | return client.admin() |
238 | .indices() |
239 | .exists(indicesExistsRequest(indexName)).actionGet().exists(); |
240 | } |
241 | |
242 | private boolean createIndex(String indexName) { |
243 | return client.admin().indices().create(Requests.createIndexRequest(indexName). |
244 | settings(new MapBuilder<String, String>().put("index.refresh_interval", "-1").map())).actionGet().acknowledged(); |
245 | } |
246 | |
247 | private <T> SearchRequestBuilder prepareSearch(Query query, Class<T> clazz){ |
248 | if(query.getIndices().isEmpty()){ |
249 | query.addIndices(retrieveIndexNameFromPersistentEntity(clazz)); |
250 | } |
251 | if(query.getTypes().isEmpty()){ |
252 | query.addTypes(retrieveTypeFromPersistentEntity(clazz)); |
253 | } |
254 | return prepareSearch(query); |
255 | } |
256 | |
257 | private SearchRequestBuilder prepareSearch(Query query){ |
258 | Assert.notNull(query.getIndices(), "No index defined for Query"); |
259 | Assert.notNull(query.getTypes(), "No type define for Query"); |
260 | |
261 | int startRecord = 0; |
262 | SearchRequestBuilder searchRequestBuilder = client.prepareSearch(toArray(query.getIndices())) |
263 | .setSearchType(DFS_QUERY_THEN_FETCH) |
264 | .setTypes(toArray(query.getTypes())); |
265 | |
266 | if(query.getPageable() != null){ |
267 | startRecord = query.getPageable().getPageNumber() * query.getPageable().getPageSize(); |
268 | searchRequestBuilder.setSize(query.getPageable().getPageSize()); |
269 | } |
270 | searchRequestBuilder.setFrom(startRecord); |
271 | |
272 | |
273 | if(!query.getFields().isEmpty()){ |
274 | searchRequestBuilder.addFields(toArray(query.getFields())); |
275 | } |
276 | |
277 | if(query.getSort() != null){ |
278 | for(Sort.Order order : query.getSort()){ |
279 | searchRequestBuilder.addSort(order.getProperty(), order.getDirection() == Sort.Direction.DESC? SortOrder.DESC : SortOrder.ASC); |
280 | } |
281 | } |
282 | return searchRequestBuilder; |
283 | } |
284 | |
285 | private IndexRequestBuilder prepareIndex(IndexQuery query){ |
286 | try { |
287 | String indexName = isBlank(query.getIndexName())? |
288 | retrieveIndexNameFromPersistentEntity(query.getObject().getClass())[0] : query.getIndexName(); |
289 | String type = isBlank(query.getType())? |
290 | retrieveTypeFromPersistentEntity(query.getObject().getClass())[0] : query.getType(); |
291 | |
292 | IndexRequestBuilder indexRequestBuilder = null; |
293 | |
294 | if(query.getId() != null){ |
295 | indexRequestBuilder = client.prepareIndex(indexName,type,query.getId()); |
296 | }else { |
297 | indexRequestBuilder = client.prepareIndex(indexName,type); |
298 | } |
299 | |
300 | indexRequestBuilder.setSource(objectMapper.writeValueAsString(query.getObject())); |
301 | |
302 | if(query.getVersion() != null){ |
303 | indexRequestBuilder.setVersion(query.getVersion()); |
304 | indexRequestBuilder.setVersionType(EXTERNAL); |
305 | } |
306 | return indexRequestBuilder; |
307 | } catch (IOException e) { |
308 | throw new ElasticsearchException("failed to index the document [id: " + query.getId() +"]",e); |
309 | } |
310 | } |
311 | |
312 | public void refresh(String indexName, boolean waitForOperation) { |
313 | client.admin().indices() |
314 | .refresh(refreshRequest(indexName).waitForOperations(waitForOperation)).actionGet(); |
315 | } |
316 | |
317 | public <T> void refresh(Class<T> clazz, boolean waitForOperation) { |
318 | ElasticsearchPersistentEntity persistentEntity = getPersistentEntityFor(clazz); |
319 | client.admin().indices() |
320 | .refresh(refreshRequest(persistentEntity.getIndexName()).waitForOperations(waitForOperation)).actionGet(); |
321 | } |
322 | |
323 | @Override |
324 | public String scan(SearchQuery query, long scrollTimeInMillis, boolean noFields) { |
325 | Assert.notNull(query.getIndices(), "No index defined for Query"); |
326 | Assert.notNull(query.getTypes(), "No type define for Query"); |
327 | Assert.notNull(query.getPageable(), "Query.pageable is required for scan & scroll"); |
328 | |
329 | SearchRequestBuilder requestBuilder = client.prepareSearch(toArray(query.getIndices())) |
330 | .setSearchType(SCAN) |
331 | .setQuery(query.getElasticsearchQuery()) |
332 | .setTypes(toArray(query.getTypes())) |
333 | .setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)) |
334 | .setFrom(0) |
335 | .setSize(query.getPageable().getPageSize()); |
336 | |
337 | if(query.getElasticsearchFilter() != null){ |
338 | requestBuilder.setFilter(query.getElasticsearchFilter()); |
339 | } |
340 | |
341 | if(noFields){ |
342 | requestBuilder.setNoFields(); |
343 | } |
344 | return requestBuilder.execute().actionGet().getScrollId(); |
345 | } |
346 | |
347 | @Override |
348 | public <T> Page<T> scroll(String scrollId, long scrollTimeInMillis, ResultsMapper<T> resultsMapper) { |
349 | SearchResponse response = client.prepareSearchScroll(scrollId) |
350 | .setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)) |
351 | .execute().actionGet(); |
352 | return resultsMapper.mapResults(response); |
353 | } |
354 | |
355 | private ElasticsearchPersistentEntity getPersistentEntityFor(Class clazz){ |
356 | Assert.isTrue(clazz.isAnnotationPresent(Document.class), "Unable to identify index name. " + |
357 | clazz.getSimpleName() + " is not a Document. Make sure the document class is annotated with @Document(indexName=\"foo\")"); |
358 | return elasticsearchConverter.getMappingContext().getPersistentEntity(clazz); |
359 | } |
360 | |
361 | private String[] retrieveIndexNameFromPersistentEntity(Class clazz){ |
362 | return new String[]{getPersistentEntityFor(clazz).getIndexName()}; |
363 | } |
364 | |
365 | private String[] retrieveTypeFromPersistentEntity(Class clazz){ |
366 | return new String[]{getPersistentEntityFor(clazz).getIndexType()}; |
367 | } |
368 | |
369 | private <T> Page<T> mapResults(SearchResponse response, final Class<T> elementType,final Pageable pageable){ |
370 | ResultsMapper<T> resultsMapper = new ResultsMapper<T>(){ |
371 | @Override |
372 | public Page<T> mapResults(SearchResponse response) { |
373 | long totalHits = response.getHits().totalHits(); |
374 | List<T> results = new ArrayList<T>(); |
375 | for (SearchHit hit : response.getHits()) { |
376 | if (hit != null) { |
377 | results.add(mapResult(hit.sourceAsString(), elementType)); |
378 | } |
379 | } |
380 | return new PageImpl<T>(results, pageable, totalHits); |
381 | } |
382 | }; |
383 | return resultsMapper.mapResults(response); |
384 | } |
385 | |
386 | private List<String> extractIds(SearchResponse response){ |
387 | List<String> ids = new ArrayList<String>(); |
388 | for (SearchHit hit : response.getHits()) { |
389 | if (hit != null) { |
390 | ids.add(hit.getId()); |
391 | } |
392 | } |
393 | return ids; |
394 | } |
395 | |
396 | private <T> T mapResult(String source, Class<T> clazz){ |
397 | if(isBlank(source)){ |
398 | return null; |
399 | } |
400 | try { |
401 | return objectMapper.readValue(source, clazz); |
402 | } catch (IOException e) { |
403 | throw new ElasticsearchException("failed to map source [ " + source + "] to class " + clazz.getSimpleName() , e); |
404 | } |
405 | } |
406 | |
407 | private static String[] toArray(List<String> values){ |
408 | String[] valuesAsArray = new String[values.size()]; |
409 | return values.toArray(valuesAsArray); |
410 | |
411 | } |
412 | |
413 | } |
414 | |