org.immutables.criteria.elasticsearch.Json Maven / Gradle / Ivy
/*
* Copyright 2019 Immutables Authors and Contributors
* Copyright 2016-2018 Apache Software Foundation (ASF)
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.immutables.criteria.elasticsearch;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.NullNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Deque;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.StreamSupport;
/**
* Mapping classes for ElasticSearch results.
*
* Some parts of this class have been copied from Apache Calcite project.
*/
class Json {
/**
* Reply of search-count
* endpoint aka {@code _count}.
*/
@JsonIgnoreProperties(ignoreUnknown = true)
static class Count {
private final long count;
@JsonCreator
Count(@JsonProperty("count") long count) {
this.count = count;
}
public long count() {
return count;
}
}
/**
* Response from Elastic
*/
@JsonIgnoreProperties(ignoreUnknown = true)
static class Result {
private final SearchHits hits;
private final Aggregations aggregations;
private final String scrollId;
private final long took;
/**
* Constructor for this instance.
* @param hits list of matched documents
* @param took time taken (in took) for this query to execute
*/
@JsonCreator
Result(@JsonProperty("hits") SearchHits hits,
@JsonProperty("aggregations") Aggregations aggregations,
@JsonProperty("_scroll_id") String scrollId,
@JsonProperty("took") long took) {
this.hits = Objects.requireNonNull(hits, "hits");
this.aggregations = aggregations;
this.scrollId = scrollId;
this.took = took;
}
SearchHits searchHits() {
return hits;
}
boolean isEmpty() {
return hits.hits().isEmpty();
}
Aggregations aggregations() {
return aggregations;
}
Duration took() {
return Duration.ofMillis(took);
}
Optional scrollId() {
return Optional.ofNullable(scrollId);
}
}
/**
* Visits leaves of the aggregation where all values are stored.
*/
static void visitValueNodes(Aggregations aggregations, Consumer> consumer) {
Objects.requireNonNull(aggregations, "aggregations");
Objects.requireNonNull(consumer, "consumer");
Map> rows = new LinkedHashMap<>();
BiConsumer cons = (r, v) ->
rows.computeIfAbsent(r, ignore -> new ArrayList<>()).add(v);
aggregations.forEach(a -> visitValueNodes(a, new ArrayList<>(), cons));
rows.forEach((k, v) -> {
if (v.stream().allMatch(val -> val instanceof GroupValue)) {
v.forEach(tuple -> {
Map groupRow = new LinkedHashMap<>(k.keys);
groupRow.put(tuple.getName(), tuple.value());
consumer.accept(groupRow);
});
} else {
Map row = new LinkedHashMap<>(k.keys);
v.forEach(val -> row.put(val.getName(), val.value()));
consumer.accept(row);
}
});
}
/**
* Visits Elasticsearch
* mapping
* properties and calls consumer for each {@code field / type} pair.
* Nested fields are represented as {@code foo.bar.qux}.
*/
static void visitMappingProperties(ObjectNode mapping,
BiConsumer consumer) {
Objects.requireNonNull(mapping, "mapping");
Objects.requireNonNull(consumer, "consumer");
visitMappingProperties(new ArrayDeque<>(), mapping, consumer);
}
private static void visitMappingProperties(Deque path,
ObjectNode mapping, BiConsumer consumer) {
Objects.requireNonNull(mapping, "mapping");
if (mapping.isMissingNode()) {
return;
}
if (mapping.has("properties")) {
// recurse
visitMappingProperties(path, (ObjectNode) mapping.get("properties"), consumer);
return;
}
if (mapping.has("type")) {
// this is leaf (register field / type mapping)
consumer.accept(String.join(".", path), mapping.get("type").asText());
return;
}
// otherwise continue visiting mapping(s)
Iterable> iter = mapping::fields;
for (Map.Entry entry : iter) {
final String name = entry.getKey();
final ObjectNode node = (ObjectNode) entry.getValue();
path.add(name);
visitMappingProperties(path, node, consumer);
path.removeLast();
}
}
/**
* Identifies a calcite row (as in relational algebra)
*/
private static class RowKey {
private final Map keys;
private final int hashCode;
private RowKey(final Map keys) {
this.keys = Objects.requireNonNull(keys, "keys");
this.hashCode = Objects.hashCode(keys);
}
private RowKey(List buckets) {
this(toMap(buckets));
}
private static Map toMap(Iterable buckets) {
return StreamSupport.stream(buckets.spliterator(), false)
.collect(LinkedHashMap::new,
(m, v) -> m.put(v.getName(), v.key()),
LinkedHashMap::putAll);
}
@Override public boolean equals(final Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
final RowKey rowKey = (RowKey) o;
return hashCode == rowKey.hashCode
&& Objects.equals(keys, rowKey.keys);
}
@Override public int hashCode() {
return this.hashCode;
}
}
private static void visitValueNodes(Aggregation aggregation, List parents,
BiConsumer consumer) {
if (aggregation instanceof MultiValue) {
// this is a leaf. publish value of the row.
RowKey key = new RowKey(parents);
consumer.accept(key, (MultiValue) aggregation);
return;
}
if (aggregation instanceof Bucket) {
Bucket bucket = (Bucket) aggregation;
if (bucket.hasNoAggregations()) {
// bucket with no aggregations is also considered a leaf node
visitValueNodes(GroupValue.of(bucket.getName(), bucket.key()), parents, consumer);
return;
}
parents.add(bucket);
bucket.getAggregations().forEach(a -> visitValueNodes(a, parents, consumer));
parents.remove(parents.size() - 1);
} else if (aggregation instanceof HasAggregations) {
HasAggregations children = (HasAggregations) aggregation;
children.getAggregations().forEach(a -> visitValueNodes(a, parents, consumer));
} else if (aggregation instanceof MultiBucketsAggregation) {
MultiBucketsAggregation multi = (MultiBucketsAggregation) aggregation;
multi.buckets().forEach(b -> visitValueNodes(b, parents, consumer));
}
}
/**
* Similar to {@code SearchHits} in ES. Container for {@link SearchHit}
*/
@JsonIgnoreProperties(ignoreUnknown = true)
static class SearchHits {
private final SearchTotal total;
private final List hits;
@JsonCreator
SearchHits(@JsonProperty("total")final SearchTotal total,
@JsonProperty("hits") final List hits) {
this.total = total;
this.hits = Objects.requireNonNull(hits, "hits");
}
public List hits() {
return this.hits;
}
public SearchTotal total() {
return total;
}
}
/**
* Container for total hits
*/
@JsonDeserialize(using = SearchTotalDeserializer.class)
static class SearchTotal {
private final long value;
SearchTotal(final long value) {
this.value = value;
}
public long value() {
return value;
}
}
/**
* Allows to de-serialize total hits structures.
*/
static class SearchTotalDeserializer extends StdDeserializer {
SearchTotalDeserializer() {
super(SearchTotal.class);
}
@Override public SearchTotal deserialize(final JsonParser parser,
final DeserializationContext ctxt)
throws IOException {
JsonNode node = parser.getCodec().readTree(parser);
return parseSearchTotal(node);
}
private static SearchTotal parseSearchTotal(JsonNode node) {
final Number value;
if (node.isNumber()) {
value = node.numberValue();
} else {
value = node.get("value").numberValue();
}
return new SearchTotal(value.longValue());
}
}
/**
* Concrete result record which matched the query. Similar to {@code SearchHit} in ES.
*/
@JsonIgnoreProperties(ignoreUnknown = true)
static class SearchHit {
/**
* ID of the document (not available in aggregations)
*/
private final String id;
private final ObjectNode source;
private final ObjectNode fields;
@JsonCreator
SearchHit(@JsonProperty("_id") final String id,
@JsonProperty("_source") final ObjectNode source,
@JsonProperty("fields") final ObjectNode fields) {
this.id = Objects.requireNonNull(id, "id");
// both can't be null
if (source == null && fields == null) {
final String message = String.format(Locale.ROOT,
"Both '_source' and 'fields' are missing for %s", id);
throw new IllegalArgumentException(message);
}
// both can't be non-null
if (source != null && fields != null) {
final String message = String.format(Locale.ROOT,
"Both '_source' and 'fields' are populated (non-null) for %s", id);
throw new IllegalArgumentException(message);
}
this.source = source;
this.fields = fields;
}
/**
* Returns id of this hit (usually document id)
* @return unique id
*/
public String id() {
return id;
}
JsonNode valueOrNull(String name) {
Objects.requireNonNull(name, "name");
if (fields != null && fields.has(name)) {
JsonNode field = fields.get(name);
if (field.isArray()) {
// return first element (or null)
Iterator iter = field.elements();
return iter.hasNext() ? iter.next() : null;
}
return field;
}
JsonNode found = source().at(name);
return found.isMissingNode() ? NullNode.getInstance() : found;
}
ObjectNode source() {
return source;
}
ObjectNode fields() {
return fields;
}
ObjectNode sourceOrFields() {
return source != null ? source : fields;
}
}
/**
* {@link Aggregation} container.
*/
@JsonDeserialize(using = AggregationsDeserializer.class)
static class Aggregations implements Iterable {
private final List aggregations;
private Map aggregationsAsMap;
Aggregations(List aggregations) {
this.aggregations = Objects.requireNonNull(aggregations, "aggregations");
}
/**
* Iterates over the {@link Aggregation}s.
*/
@Override public final Iterator iterator() {
return asList().iterator();
}
/**
* The list of {@link Aggregation}s.
*/
final List asList() {
return Collections.unmodifiableList(aggregations);
}
/**
* Returns the {@link Aggregation}s keyed by aggregation name. Lazy init.
*/
final Map asMap() {
if (aggregationsAsMap == null) {
Map map = new LinkedHashMap<>(aggregations.size());
for (Aggregation aggregation : aggregations) {
map.put(aggregation.getName(), aggregation);
}
this.aggregationsAsMap = Collections.unmodifiableMap(map);
}
return aggregationsAsMap;
}
/**
* Returns the aggregation that is associated with the specified name.
*/
@SuppressWarnings("unchecked")
public final A get(String name) {
return (A) asMap().get(name);
}
@Override public final boolean equals(Object obj) {
if (obj == null || getClass() != obj.getClass()) {
return false;
}
return aggregations.equals(((Aggregations) obj).aggregations);
}
@Override public final int hashCode() {
return Objects.hash(getClass(), aggregations);
}
}
/**
* Identifies all aggregations
*/
interface Aggregation {
/**
* @return The name of this aggregation.
*/
String getName();
}
/**
* Allows traversing aggregations tree
*/
interface HasAggregations {
Aggregations getAggregations();
}
/**
* An aggregation that returns multiple buckets
*/
static class MultiBucketsAggregation implements Aggregation {
private final String name;
private final List buckets;
MultiBucketsAggregation(final String name,
final List buckets) {
this.name = name;
this.buckets = buckets;
}
/**
* @return The buckets of this aggregation.
*/
List buckets() {
return buckets;
}
@Override public String getName() {
return name;
}
}
/**
* A bucket represents a criteria to which all documents that fall in it adhere to.
* It is also uniquely identified
* by a key, and can potentially hold sub-aggregations computed over all documents in it.
*/
static class Bucket implements HasAggregations, Aggregation {
private final Object key;
private final String name;
private final Aggregations aggregations;
Bucket(final Object key,
final String name,
final Aggregations aggregations) {
this.key = key; // key can be set after construction
this.name = Objects.requireNonNull(name, "name");
this.aggregations = Objects.requireNonNull(aggregations, "aggregations");
}
/**
* @return The key associated with the bucket
*/
Object key() {
return key;
}
/**
* @return The key associated with the bucket as a string
*/
String keyAsString() {
return Objects.toString(key());
}
/**
* Means current bucket has no aggregations.
*/
boolean hasNoAggregations() {
return aggregations.asList().isEmpty();
}
/**
* @return The sub-aggregations of this bucket
*/
@Override public Aggregations getAggregations() {
return aggregations;
}
@Override public String getName() {
return name;
}
}
/**
* Multi value aggregatoin like
* Stats
*/
static class MultiValue implements Aggregation {
private final String name;
private final Map values;
MultiValue(final String name, final Map values) {
this.name = Objects.requireNonNull(name, "name");
this.values = Objects.requireNonNull(values, "values");
}
@Override public String getName() {
return name;
}
Map values() {
return values;
}
/**
* For single value. Returns single value represented by this leaf aggregation.
* @return value corresponding to {@code value}
*/
Object value() {
if (!values().containsKey("value")) {
String message = String.format(Locale.ROOT, "'value' field not present in "
+ "%s aggregation", getName());
throw new IllegalStateException(message);
}
return values().get("value");
}
}
/**
* Distinguishes from {@link MultiValue}.
* In order that rows which have the same key can be put into result map.
*/
static class GroupValue extends MultiValue {
GroupValue(String name, Map values) {
super(name, values);
}
/**
* Constructs a {@link GroupValue} instance with a single value.
*/
static GroupValue of(String name, Object value) {
return new GroupValue(name, Collections.singletonMap("value", value));
}
}
/**
* Allows to de-serialize nested aggregation structures.
*/
static class AggregationsDeserializer extends StdDeserializer {
private static final Set IGNORE_TOKENS =
ImmutableSet.of("meta", "buckets", "value", "values", "value_as_string",
"doc_count", "key", "key_as_string");
AggregationsDeserializer() {
super(Aggregations.class);
}
@Override public Aggregations deserialize(final JsonParser parser,
final DeserializationContext ctxt)
throws IOException {
ObjectNode node = parser.getCodec().readTree(parser);
return parseAggregations(parser, node);
}
private static Aggregations parseAggregations(JsonParser parser, ObjectNode node)
throws JsonProcessingException {
List aggregations = new ArrayList<>();
Iterable> iter = node::fields;
for (Map.Entry entry : iter) {
final String name = entry.getKey();
final JsonNode value = entry.getValue();
Aggregation agg = null;
if (value.has("buckets")) {
agg = parseBuckets(parser, name, (ArrayNode) value.get("buckets"));
} else if (value.isObject() && !IGNORE_TOKENS.contains(name)) {
// leaf
agg = parseValue(parser, name, (ObjectNode) value);
}
if (agg != null) {
aggregations.add(agg);
}
}
return new Aggregations(aggregations);
}
private static MultiValue parseValue(JsonParser parser, String name, ObjectNode node)
throws JsonProcessingException {
@SuppressWarnings("unchecked")
Map values = (Map) parser.getCodec().treeToValue(node, Map.class);
return new MultiValue(name, values);
}
private static Aggregation parseBuckets(JsonParser parser, String name, ArrayNode nodes)
throws JsonProcessingException {
List buckets = new ArrayList<>(nodes.size());
for (JsonNode b: nodes) {
buckets.add(parseBucket(parser, name, (ObjectNode) b));
}
return new MultiBucketsAggregation(name, buckets);
}
/**
* Determines if current key is a missing field key. Missing key is returned when document
* does not have pivoting attribute (example {@code GROUP BY _MAP['a.b.missing']}). It helps
* grouping documents which don't have a field. In relational algebra this
* would normally be {@code null}.
*
* Please note that missing value is different for each type.
*
* @param key current {@code key} (usually string) as returned by ES
* @return {@code true} if this value
*/
private static boolean isMissingBucket(JsonNode key) {
return Mapping.Datatype.isMissingValue(key);
}
private static Bucket parseBucket(JsonParser parser, String name, ObjectNode node)
throws JsonProcessingException {
if (!node.has("key")) {
throw new IllegalArgumentException("No 'key' attribute for " + node);
}
final JsonNode keyNode = node.get("key");
final Object key;
if (isMissingBucket(keyNode) || keyNode.isNull()) {
key = null;
} else if (keyNode.isTextual()) {
key = keyNode.textValue();
} else if (keyNode.isNumber()) {
key = keyNode.numberValue();
} else if (keyNode.isBoolean()) {
key = keyNode.booleanValue();
} else {
// don't usually expect keys to be Objects
key = parser.getCodec().treeToValue(node, Map.class);
}
return new Bucket(key, name, parseAggregations(parser, node));
}
}
}