diff --git a/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java b/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java index 3021f5b31606e..151dd8b9e61ba 100644 --- a/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java +++ b/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java @@ -178,7 +178,7 @@ public PercolateQueryBuilder(String field, String documentType, List documentSupplier) { + protected PercolateQueryBuilder(String field, String documentType, Supplier documentSupplier) { if (field == null) { throw new IllegalArgumentException("[field] is a required argument"); } @@ -491,8 +491,12 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) { if (source == null) { return this; // not executed yet } else { - return new PercolateQueryBuilder(field, documentType, Collections.singletonList(source), - XContentHelper.xContentType(source)); + PercolateQueryBuilder rewritten = new PercolateQueryBuilder(field, documentType, + Collections.singletonList(source), XContentHelper.xContentType(source)); + if (name != null) { + rewritten.setName(name); + } + return rewritten; } } GetRequest getRequest; @@ -527,7 +531,12 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) { listener.onResponse(null); }, listener::onFailure)); }); - return new PercolateQueryBuilder(field, documentType, documentSupplier::get); + + PercolateQueryBuilder rewritten = new PercolateQueryBuilder(field, documentType, documentSupplier::get); + if (name != null) { + rewritten.setName(name); + } + return rewritten; } @Override @@ -626,6 +635,10 @@ XContentType getXContentType() { return documentXContentType; } + public String getQueryName() { + return name; + } + static IndexSearcher createMultiDocumentSearcher(Analyzer analyzer, Collection docs) { RAMDirectory ramDirectory = new RAMDirectory(); try (IndexWriter indexWriter = new IndexWriter(ramDirectory, new IndexWriterConfig(analyzer))) { diff --git a/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolateQueryBuilderTests.java b/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolateQueryBuilderTests.java index 6053a92b54a20..5b4dc61090042 100644 --- a/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolateQueryBuilderTests.java +++ b/modules/percolator/src/test/java/org/elasticsearch/percolator/PercolateQueryBuilderTests.java @@ -54,6 +54,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Supplier; import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO; import static org.hamcrest.Matchers.equalTo; @@ -331,4 +332,29 @@ public void testFieldAlias() throws IOException { assertEquals(query.getVerifiedMatchesQuery(), aliasQuery.getVerifiedMatchesQuery()); } + public void testSettingNameWhileRewriting() { + String testName = "name1"; + QueryShardContext shardContext = createShardContext(); + PercolateQueryBuilder percolateQueryBuilder = doCreateTestQueryBuilder(true); + percolateQueryBuilder.setName(testName); + + QueryBuilder rewrittenQueryBuilder = percolateQueryBuilder.doRewrite(shardContext); + + assertEquals(testName, ((PercolateQueryBuilder) rewrittenQueryBuilder).getQueryName()); + assertNotEquals(rewrittenQueryBuilder, percolateQueryBuilder); + } + + public void testSettingNameWhileRewritingWhenDocumentSupplierAndSourceNotNull() { + Supplier supplier = () -> new BytesArray("{\"test\": \"test\"}"); + String testName = "name1"; + QueryShardContext shardContext = createShardContext(); + PercolateQueryBuilder percolateQueryBuilder = new PercolateQueryBuilder(queryField, null, supplier); + percolateQueryBuilder.setName(testName); + + QueryBuilder rewrittenQueryBuilder = percolateQueryBuilder.doRewrite(shardContext); + + assertEquals(testName, ((PercolateQueryBuilder) rewrittenQueryBuilder).getQueryName()); + assertNotEquals(rewrittenQueryBuilder, percolateQueryBuilder); + } + }