Migrate OpenSearchBackend and related search functionality to the new Java client

graylog-storage-opensearch3/src/main/java/org/graylog/storage/opensearch3/ViewsOSBackendModule.java

@@ -46,7 +46,6 @@
 import org.graylog.plugins.views.search.searchtypes.pivot.series.Sum;
 import org.graylog.plugins.views.search.searchtypes.pivot.series.SumOfSquares;
 import org.graylog.plugins.views.search.searchtypes.pivot.series.Variance;
-import org.graylog.shaded.opensearch2.org.opensearch.search.aggregations.Aggregation;
 import org.graylog.storage.opensearch3.views.OSGeneratedQueryContext;
 import org.graylog.storage.opensearch3.views.OpenSearchBackend;
 import org.graylog.storage.opensearch3.views.export.OpenSearchExportBackend;
@@ -130,7 +129,7 @@ private void registerPivotBucketHandler(
         pivotBucketHandlerBinder().addBinding(name).to(implementation);
     }
 
-    protected MapBinder<String, OSPivotSeriesSpecHandler<? extends SeriesSpec, ? extends Aggregation>> pivotSeriesHandlerBinder() {
+    protected MapBinder<String, OSPivotSeriesSpecHandler<? extends SeriesSpec>> pivotSeriesHandlerBinder() {
         return MapBinder.newMapBinder(binder(),
                 TypeLiteral.get(String.class),
                 new TypeLiteral<>() {});
@@ -139,7 +138,7 @@ private void registerPivotBucketHandler(
 
     private void registerPivotSeriesHandler(
             String name,
-            Class<? extends OSPivotSeriesSpecHandler<? extends SeriesSpec, ? extends Aggregation>> implementation
+            Class<? extends OSPivotSeriesSpecHandler<? extends SeriesSpec>> implementation
     ) {
         pivotSeriesHandlerBinder().addBinding(name).to(implementation);
     }

graylog-storage-opensearch3/src/main/java/org/graylog/storage/opensearch3/indextemplates/OSSerializationUtils.java

@@ -18,8 +18,11 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JsonNode;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
+import jakarta.json.Json;
+import jakarta.json.JsonWriter;
 import jakarta.json.stream.JsonParser;
 import org.opensearch.client.json.JsonData;
 import org.opensearch.client.json.JsonpDeserializer;
@@ -28,7 +31,9 @@
 import org.opensearch.client.opensearch.core.SearchRequest;
 import org.opensearch.client.opensearch.core.msearch.RequestItem;
 
+import java.io.IOException;
 import java.io.StringReader;
+import java.io.StringWriter;
 import java.util.Map;
 import java.util.stream.Collectors;
 
@@ -46,11 +51,62 @@ public OSSerializationUtils() {
         this.jsonpMapper = new JacksonJsonpMapper();
     }
 
+    public static Object valueNode(JsonNode jsonNode) {
+        if (jsonNode.isInt()) {
+            return jsonNode.asInt();
+        } else if (jsonNode.isLong()) {
+            return jsonNode.asLong();
+        } else if (jsonNode.isIntegralNumber()) {
+            return jsonNode.asLong();
+        } else if (jsonNode.isFloatingPointNumber()) {
+            return jsonNode.asDouble();
+        } else if (jsonNode.isBoolean()) {
+            return jsonNode.asBoolean();
+        } else if (jsonNode.isNull()) {
+            return null;
+        } else {
+            return jsonNode.asText();
+        }
+    }
+
+    public Object toObject(JsonData jsonData) {
+        try {
+            return valueNode(jsonpMapper.objectMapper().readTree(toJson(jsonData)));
+        } catch (JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public Map<String, Object> toMap(final JsonData openSearchSerializableObject) {
+        String json = toJson(openSearchSerializableObject);
+        return toMap(json);
+    }
 
-    public Map<String, Object> toMap(final PlainJsonSerializable openSearchSerializableObject) throws JsonProcessingException {
-        return this.jsonpMapper.objectMapper().readValue(openSearchSerializableObject.toJsonString(), new TypeReference<>() {});
+    private String toJson(final JsonData openSearchSerializableObject) {
+        try (
+                StringWriter writer = new StringWriter();
+                JsonWriter jsonWriter = Json.createWriter(writer)
+        ) {
+            jsonWriter.write(openSearchSerializableObject.toJson());
+            return writer.toString();
+        } catch (IOException e) {
+            throw new RuntimeException("Error serializing json", e);
+        }
     }
 
+    public Map<String, Object> toMap(final PlainJsonSerializable openSearchSerializableObject) {
+        return toMap(openSearchSerializableObject.toJsonString());
+    }
+
+    public Map<String, Object> toMap(final String json) {
+        try {
+            return this.jsonpMapper.objectMapper().readValue(json, new TypeReference<>() {});
+        } catch (IOException e) {
+            throw new RuntimeException("Error serializing json", e);
+        }
+    }
+
+
     public Map<String, JsonData> toJsonDataMap(final Map<String, Object> map) {
         return map.entrySet()
                 .stream()

graylog-storage-opensearch3/src/main/java/org/graylog/storage/opensearch3/views/MutableSearchRequestBuilder.java

@@ -0,0 +1,233 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program 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
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+
+package org.graylog.storage.opensearch3.views;
+
+import org.graylog.storage.opensearch3.views.searchtypes.pivot.MutableNamedAggregationBuilder;
+import org.opensearch.client.opensearch._types.ExpandWildcard;
+import org.opensearch.client.opensearch._types.FieldValue;
+import org.opensearch.client.opensearch._types.SortOptions;
+import org.opensearch.client.opensearch._types.Time;
+import org.opensearch.client.opensearch._types.query_dsl.Query;
+import org.opensearch.client.opensearch.core.SearchRequest;
+import org.opensearch.client.opensearch.core.search.Highlight;
+import org.opensearch.client.opensearch.core.search.SourceConfig;
+import org.opensearch.client.opensearch.core.search.TrackHits;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Wrapper for SearchRequest.Builder.
+ * All builders in the new client are immutable, write-only and cannot be built more than once.
+ * Since we are mutating a reference to the builder in the handlers, we need to be able to read fields
+ * that are already set (e.g. to check the type or add to the existing query). Also, we need to be able to
+ * copy an existing builder without building and using toBuilder().
+ * This class provides the structure for this.
+ */
+public class MutableSearchRequestBuilder {
+
+    Query query;
+    Integer from;
+    Integer size;
+    TrackHits trackTotalHits;
+    List<String> indices;
+    Boolean allowNoIndices;
+    Boolean ignoreUnavailable;
+    List<ExpandWildcard> expandWildcards;
+    Time cancelAfterTimeInterval;
+    String preference;
+    List<SortOptions> sort;
+    SourceConfig source;
+    Highlight highlight;
+    List<MutableNamedAggregationBuilder> aggregations = new ArrayList<>();
+    List<FieldValue> searchAfter;
+
+    public MutableSearchRequestBuilder query(Query query) {
+        this.query = query;
+        return this;
+    }
+
+    public Query query() {
+        return query;
+    }
+
+    public MutableSearchRequestBuilder from(Integer from) {
+        this.from = from;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder size(Integer size) {
+        this.size = size;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder trackTotalHits(TrackHits trackTotalHits) {
+        this.trackTotalHits = trackTotalHits;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder index(List<String> indices) {
+        this.indices = indices;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder allowNoIndices(Boolean allowNoIndices) {
+        this.allowNoIndices = allowNoIndices;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder ignoreUnavailable(Boolean ignoreUnavailable) {
+        this.ignoreUnavailable = ignoreUnavailable;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder expandWildcards(ExpandWildcard... expandWildcards) {
+        this.expandWildcards = Arrays.asList(expandWildcards);
+        return this;
+    }
+
+    public MutableSearchRequestBuilder expandWildcards(List<ExpandWildcard> expandWildcards) {
+        this.expandWildcards = expandWildcards;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder cancelAfterTimeInterval(Time cancelAfterTimeInterval) {
+        this.cancelAfterTimeInterval = cancelAfterTimeInterval;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder preference(String preference) {
+        this.preference = preference;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder sort(List<SortOptions> sort) {
+        this.sort = sort;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder sort(SortOptions... sort) {
+        if (this.sort == null) {
+            this.sort = new ArrayList<>();
+        }
+        this.sort.addAll(Arrays.asList(sort));
+        return this;
+    }
+
+    public MutableSearchRequestBuilder source(SourceConfig source) {
+        this.source = source;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder highlight(Highlight highlight) {
+        this.highlight = highlight;
+        return this;
+    }
+
+    public Highlight highlight() {
+        return highlight;
+    }
+
+    public MutableSearchRequestBuilder aggregations(List<MutableNamedAggregationBuilder> aggregations) {
+        this.aggregations = aggregations;
+        return this;
+    }
+
+    public MutableSearchRequestBuilder aggregation(MutableNamedAggregationBuilder aggregation) {
+        this.aggregations.add(aggregation);
+        return this;
+    }
+
+    public MutableSearchRequestBuilder searchAfter(List<FieldValue> searchAfter) {
+        this.searchAfter = searchAfter;
+        return this;
+    }
+
+    SearchRequest build() {
+        return SearchRequest.of(b -> {
+            if (query != null) {
+                b.query(query);
+            }
+            if (from != null) {
+                b.from(from);
+            }
+            if (size != null) {
+                b.size(size);
+            }
+            if (trackTotalHits != null) {
+                b.trackTotalHits(trackTotalHits);
+            }
+            if (indices != null) {
+                b.index(indices);
+            }
+            if (allowNoIndices != null) {
+                b.allowNoIndices(allowNoIndices);
+            }
+            if (ignoreUnavailable != null) {
+                b.ignoreUnavailable(ignoreUnavailable);
+            }
+            if (expandWildcards != null) {
+                b.expandWildcards(expandWildcards);
+            }
+            if (sort != null) {
+                b.sort(sort);
+            }
+            if (source != null) {
+                b.source(source);
+            }
+            if (highlight != null) {
+                b.highlight(highlight);
+            }
+            if (aggregations != null) {
+                aggregations.forEach(aggregation -> {
+                    b.aggregations(aggregation.getName(), aggregation.build());
+                });
+            }
+            if (searchAfter != null) {
+                b.searchAfter(searchAfter);
+            }
+            return b;
+        });
+    }
+
+    public MutableSearchRequestBuilder copy() {
+        final MutableSearchRequestBuilder copy = new MutableSearchRequestBuilder()
+                .query(query)
+                .from(from)
+                .size(size)
+                .trackTotalHits(trackTotalHits)
+                .index(indices)
+                .allowNoIndices(allowNoIndices)
+                .ignoreUnavailable(ignoreUnavailable)
+                .expandWildcards(expandWildcards)
+                .cancelAfterTimeInterval(cancelAfterTimeInterval)
+                .preference(preference)
+                .source(source)
+                .highlight(highlight)
+                .searchAfter(searchAfter);
+        copy.aggregations = new ArrayList<>(aggregations);
+        copy.sort = sort != null ? new ArrayList<>(sort) : null;
+        return copy;
+    }
+
+    @Override
+    public String toString() {
+        return build().toJsonString();
+    }
+}