-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
EsUtils.java
124 lines (109 loc) · 3.96 KB
/
EsUtils.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
/*
* SonarQube, open source software quality management tool.
* Copyright (C) 2008-2014 SonarSource
* mailto:contact AT sonarsource DOT com
*
* SonarQube is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 3 of the License, or (at your option) any later version.
*
* SonarQube is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public License
* along with this program; if not, write to the Free Software Foundation,
* Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
*/
package org.sonar.server.es;
import com.google.common.base.Function;
import com.google.common.collect.Lists;
import org.elasticsearch.action.search.SearchScrollRequestBuilder;
import org.elasticsearch.common.joda.time.format.ISODateTimeFormat;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.sonar.server.search.BaseDoc;
import javax.annotation.CheckForNull;
import javax.annotation.Nullable;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Queue;
public class EsUtils {
public static final int SCROLL_TIME_IN_MINUTES = 3;
private EsUtils() {
// only static methods
}
public static <D extends BaseDoc> List<D> convertToDocs(SearchHits hits, Function<Map<String, Object>, D> converter) {
List<D> docs = new ArrayList<>();
for (SearchHit hit : hits.getHits()) {
docs.add(converter.apply(hit.getSource()));
}
return docs;
}
public static LinkedHashMap<String, Long> termsToMap(Terms terms) {
LinkedHashMap<String, Long> map = new LinkedHashMap<>();
List<Terms.Bucket> buckets = terms.getBuckets();
for (Terms.Bucket bucket : buckets) {
map.put(bucket.getKey(), bucket.getDocCount());
}
return map;
}
public static List<String> termsKeys(Terms terms) {
return Lists.transform(terms.getBuckets(), new Function<Terms.Bucket, String>() {
@Override
public String apply(Terms.Bucket bucket) {
return bucket.getKey();
}
});
}
@CheckForNull
public static Date parseDateTime(@Nullable String s) {
if (s != null) {
return ISODateTimeFormat.dateTime().parseDateTime(s).toDate();
}
return null;
}
@CheckForNull
public static String formatDateTime(@Nullable Date date) {
if (date != null) {
return ISODateTimeFormat.dateTime().print(date.getTime());
}
return null;
}
public static <D extends BaseDoc> Iterator<D> scroll(final EsClient esClient, final String scrollId, final Function<Map<String, Object>, D> docConverter) {
return new Iterator<D>() {
private final Queue<SearchHit> hits = new ArrayDeque<>();
@Override
public boolean hasNext() {
if (hits.isEmpty()) {
SearchScrollRequestBuilder esRequest = esClient.prepareSearchScroll(scrollId)
.setScroll(TimeValue.timeValueMinutes(SCROLL_TIME_IN_MINUTES));
Collections.addAll(hits, esRequest.get().getHits().getHits());
}
return !hits.isEmpty();
}
@Override
public D next() {
if (!hasNext()) {
throw new NoSuchElementException();
}
return docConverter.apply(hits.poll().getSource());
}
@Override
public void remove() {
throw new UnsupportedOperationException("Cannot remove item when scrolling");
}
};
}
}