Add a test module to simulate errors and warnings in search requests

test/external-modules/error-query/build.gradle

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+esplugin {
+  description 'A test module that exposes a way to simulate search shard failures and warnings'
+  classname 'org.elasticsearch.search.query.ErrorQueryPlugin'
+}
+
+restResources {
+  restApi {
+    include '_common', 'indices', 'index', 'cluster', 'search'
+  }
+}

test/external-modules/error-query/src/main/java/org/elasticsearch/search/query/ErrorQueryBuilder.java

@@ -0,0 +1,130 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.query;
+
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.logging.HeaderWarning;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
+
+/**
+ * A test query that can simulate errors and warnings when executing a shard request.
+ */
+public class ErrorQueryBuilder extends AbstractQueryBuilder<ErrorQueryBuilder> {
+    public static final String NAME = "error_query";
+
+    private List<IndexError> indices;
+
+    public ErrorQueryBuilder(List<IndexError> indices) {
+        this.indices = Objects.requireNonNull(indices);
+    }
+
+    public ErrorQueryBuilder(StreamInput in) throws IOException {
+        super(in);
+        this.indices = in.readList(IndexError::new);
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeList(indices);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
+        // Disable the request cache
+        context.nowInMillis();
+
+        IndexError error = null;
+        for (IndexError index : indices) {
+            if (context.indexMatches(index.getIndexName())) {
+                error = index;
+                break;
+            }
+        }
+        if (error == null) {
+            return new MatchAllDocsQuery();
+        }
+        if (error.getShardIds() != null) {
+            boolean match = false;
+            for (int shardId : error.getShardIds()) {
+                if (context.getShardId() == shardId) {
+                    match = true;
+                    break;
+                }
+            }
+            if (match == false) {
+                return new MatchAllDocsQuery();
+            }
+        }
+        final String header = "[" + context.index().getName() + "][" + context.getShardId() + "]";
+        if (error.getErrorType() == IndexError.ERROR_TYPE.WARNING) {
+            HeaderWarning.addWarning(header + " " + error.getMessage());
+            return new MatchAllDocsQuery();
+        } else {
+            throw new RuntimeException(header + " " + error.getMessage());
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    static final ConstructingObjectParser<ErrorQueryBuilder, String> PARSER = new ConstructingObjectParser<>(NAME, false, (args, name) -> {
+        ErrorQueryBuilder q = new ErrorQueryBuilder((List<IndexError>) args[0]);
+        final float boost = args[1] == null ? DEFAULT_BOOST : (Float) args[1];
+        final String queryName = (String) args[2];
+        q.boost = boost;
+        q.queryName = queryName;
+        return q;
+    });
+
+    static {
+        PARSER.declareObjectArray(constructorArg(), (p, c) -> IndexError.PARSER.parse(p, c), new ParseField("indices"));
+        PARSER.declareFloat(ConstructingObjectParser.optionalConstructorArg(), BOOST_FIELD);
+        PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), NAME_FIELD);
+    }
+
+    @Override
+    protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(NAME);
+        builder.startArray("indices");
+        for (IndexError indexError : indices) {
+            builder.startObject();
+            indexError.toXContent(builder, params);
+            builder.endObject();
+        }
+        builder.endArray();
+        printBoostAndQueryName(builder);
+        builder.endObject();
+    }
+
+    @Override
+    protected boolean doEquals(ErrorQueryBuilder other) {
+        return Objects.equals(indices, other.indices);
+    }
+
+    @Override
+    protected int doHashCode() {
+        return Objects.hash(indices);
+    }
+}

test/external-modules/error-query/src/main/java/org/elasticsearch/search/query/ErrorQueryPlugin.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.query;
+
+import java.util.List;
+
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.SearchPlugin;
+
+import static java.util.Collections.singletonList;
+
+/**
+ * Test plugin that exposes a way to simulate search shard failures and warnings.
+ */
+public class ErrorQueryPlugin extends Plugin implements SearchPlugin {
+    public ErrorQueryPlugin() {}
+
+    @Override
+    public List<QuerySpec<?>> getQueries() {
+        return singletonList(new QuerySpec<>(ErrorQueryBuilder.NAME, ErrorQueryBuilder::new, p -> ErrorQueryBuilder.PARSER.parse(p, null)));
+    }
+}

test/external-modules/error-query/src/main/java/org/elasticsearch/search/query/IndexError.java

@@ -0,0 +1,133 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.query;
+
+import org.elasticsearch.common.Nullable;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ToXContentFragment;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Locale;
+import java.util.Objects;
+
+import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
+import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
+
+public class IndexError implements Writeable, ToXContentFragment {
+    enum ERROR_TYPE {
+        WARNING,
+        EXCEPTION
+    }
+
+    private final String indexName;
+    private final int[] shardIds;
+    private final ERROR_TYPE errorType;
+    private final String message;
+
+    public IndexError(String indexName, int[] shardIds, ERROR_TYPE errorType, String message) {
+        this.indexName = indexName;
+        this.shardIds = shardIds;
+        this.errorType = errorType;
+        this.message = message;
+    }
+
+    public IndexError(StreamInput in) throws IOException {
+        this.indexName = in.readString();
+        this.shardIds = in.readBoolean() ? in.readIntArray() : null;
+        this.errorType = in.readEnum(ERROR_TYPE.class);
+        this.message = in.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(indexName);
+        out.writeBoolean(shardIds != null);
+        if (shardIds != null) {
+            out.writeIntArray(shardIds);
+        }
+        out.writeEnum(errorType);
+        out.writeString(message);
+    }
+
+    public String getIndexName() {
+        return indexName;
+    }
+
+    @Nullable
+    public int[] getShardIds() {
+        return shardIds;
+    }
+
+    public ERROR_TYPE getErrorType() {
+        return errorType;
+    }
+
+    public String getMessage() {
+        return message;
+    }
+
+    @SuppressWarnings("unchecked")
+    static final ConstructingObjectParser<IndexError, String> PARSER = new ConstructingObjectParser<>(
+        "index_error",
+        false,
+        (args, name) -> {
+            List<Integer> lst = (List<Integer>) args[1];
+            int[] shardIds = lst == null ? null : lst.stream().mapToInt(i -> i).toArray();
+            return new IndexError(
+                (String) args[0],
+                shardIds,
+                ERROR_TYPE.valueOf(((String) args[2]).toUpperCase(Locale.ROOT)),
+                (String) args[3]
+            );
+        }
+    );
+
+    static {
+        PARSER.declareString(constructorArg(), new ParseField("name"));
+        PARSER.declareIntArray(optionalConstructorArg(), new ParseField("shard_ids"));
+        PARSER.declareString(constructorArg(), new ParseField("error_type"));
+        PARSER.declareString(constructorArg(), new ParseField("message"));
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field("name", indexName);
+        if (shardIds != null) {
+            builder.field("shard_ids", shardIds);
+        }
+        builder.field("error_type", errorType.toString());
+        builder.field("message", message);
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        IndexError that = (IndexError) o;
+        return indexName.equals(that.indexName)
+            && Arrays.equals(shardIds, that.shardIds)
+            && errorType == that.errorType
+            && message.equals(that.message);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = Objects.hash(indexName, errorType, message);
+        result = 31 * result + Arrays.hashCode(shardIds);
+        return result;
+    }
+}

test/external-modules/error-query/src/test/java/org/elasticsearch/search/query/ErrorQueryBuilderTests.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.search.query;
+
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.AbstractQueryTestCase;
+import org.elasticsearch.test.TestGeoShapeFieldMapperPlugin;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+
+public class ErrorQueryBuilderTests extends AbstractQueryTestCase<ErrorQueryBuilder> {
+    @Override
+    protected Collection<Class<? extends Plugin>> getPlugins() {
+        return Arrays.asList(ErrorQueryPlugin.class, TestGeoShapeFieldMapperPlugin.class);
+    }
+
+    @Override
+    protected ErrorQueryBuilder doCreateTestQueryBuilder() {
+        int numIndex = randomIntBetween(0, 5);
+        List<IndexError> indices = new ArrayList<>();
+        for (int i = 0; i < numIndex; i++) {
+            String indexName = randomAlphaOfLengthBetween(5, 30);
+            int numShards = randomIntBetween(0, 3);
+            int[] shardIds = numShards > 0 ? new int[numShards] : null;
+            for (int j = 0; j < numShards; j++) {
+                shardIds[j] = j;
+            }
+            indices.add(
+                new IndexError(indexName, shardIds, randomFrom(IndexError.ERROR_TYPE.values()), randomAlphaOfLengthBetween(5, 100))
+            );
+        }
+
+        return new ErrorQueryBuilder(indices);
+    }
+
+    @Override
+    protected void doAssertLuceneQuery(ErrorQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
+        assertEquals(new MatchAllDocsQuery(), query);
+    }
+
+    @Override
+    public void testCacheability() throws IOException {
+        ErrorQueryBuilder queryBuilder = createTestQueryBuilder();
+        SearchExecutionContext context = createSearchExecutionContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new SearchExecutionContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
+    }
+}