| 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 | |