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 44200823b6d66..67db4ba9c0e0c 100644 --- a/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java +++ b/modules/percolator/src/main/java/org/elasticsearch/percolator/PercolateQueryBuilder.java @@ -181,7 +181,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"); } @@ -519,8 +519,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; @@ -555,7 +559,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 @@ -654,6 +663,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); + } + }