Allow terms query in _rollup_search (#30973)

This change adds the `terms` query to the list of accepted queries
for the _rollup_search endpoint.

x-pack/plugin/rollup/src/main/java/org/elasticsearch/xpack/rollup/action/TransportRollupSearchAction.java

@@ -39,6 +39,7 @@ import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.RangeQueryBuilder;
 import org.elasticsearch.index.query.TermQueryBuilder;
+import org.elasticsearch.index.query.TermsQueryBuilder;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -66,6 +67,7 @@ import org.joda.time.DateTimeZone;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
@@ -271,16 +273,52 @@ public class TransportRollupSearchAction extends TransportAction<SearchRequest,
                     rewriteQuery(((BoostingQueryBuilder)builder).positiveQuery(), jobCaps));
         } else if (builder.getWriteableName().equals(DisMaxQueryBuilder.NAME)) {
             DisMaxQueryBuilder rewritten = new DisMaxQueryBuilder();
-            ((DisMaxQueryBuilder)builder).innerQueries().forEach(query -> rewritten.add(rewriteQuery(query, jobCaps)));
+            ((DisMaxQueryBuilder) builder).innerQueries().forEach(query -> rewritten.add(rewriteQuery(query, jobCaps)));
             return rewritten;
-        } else if (builder.getWriteableName().equals(RangeQueryBuilder.NAME) || builder.getWriteableName().equals(TermQueryBuilder.NAME)) {
+        } else if (builder.getWriteableName().equals(RangeQueryBuilder.NAME)) {
+            RangeQueryBuilder range = (RangeQueryBuilder) builder;
+            String fieldName = range.fieldName();
+            // Many range queries don't include the timezone because the default is UTC, but the query
+            // builder will return null so we need to set it here
+            String timeZone = range.timeZone() == null ? DateTimeZone.UTC.toString() : range.timeZone();
 
-            String fieldName = builder.getWriteableName().equals(RangeQueryBuilder.NAME)
+            String rewrittenFieldName = rewriteFieldName(jobCaps, RangeQueryBuilder.NAME, fieldName, timeZone);
-                    ? ((RangeQueryBuilder)builder).fieldName()
+            RangeQueryBuilder rewritten = new RangeQueryBuilder(rewrittenFieldName)
-                    : ((TermQueryBuilder)builder).fieldName();
+                .from(range.from())
+                .to(range.to())
+                .includeLower(range.includeLower())
+                .includeUpper(range.includeUpper());
+            if (range.timeZone() != null) {
+                rewritten.timeZone(range.timeZone());
+            }
+            if (range.format() != null) {
+                rewritten.format(range.format());
+            }
+            return rewritten;
+        } else if (builder.getWriteableName().equals(TermQueryBuilder.NAME)) {
+            TermQueryBuilder term = (TermQueryBuilder) builder;
+            String fieldName = term.fieldName();
+            String rewrittenFieldName =  rewriteFieldName(jobCaps, TermQueryBuilder.NAME, fieldName, null);
+            return new TermQueryBuilder(rewrittenFieldName, term.value());
+        } else if (builder.getWriteableName().equals(TermsQueryBuilder.NAME)) {
+            TermsQueryBuilder terms = (TermsQueryBuilder) builder;
+            String fieldName = terms.fieldName();
+            String rewrittenFieldName =  rewriteFieldName(jobCaps, TermQueryBuilder.NAME, fieldName, null);
+            return new TermsQueryBuilder(rewrittenFieldName, terms.values());
+        } else if (builder.getWriteableName().equals(MatchAllQueryBuilder.NAME)) {
+            // no-op
+            return builder;
+        } else {
+            throw new IllegalArgumentException("Unsupported Query in search request: [" + builder.getWriteableName() + "]");
+        }
+    }
 
-            List<String> incorrectTimeZones = new ArrayList<>();
+    private static String rewriteFieldName(Set<RollupJobCaps> jobCaps,
-            List<String> rewrittenFieldName = jobCaps.stream()
+                                           String builderName,
+                                           String fieldName,
+                                           String timeZone) {
+        List<String> incompatibleTimeZones = timeZone == null ? Collections.emptyList() : new ArrayList<>();
+        List<String> rewrittenFieldNames = jobCaps.stream()
             // We only care about job caps that have the query's target field
             .filter(caps -> caps.getFieldCaps().keySet().contains(fieldName))
             .map(caps -> {
@@ -291,18 +329,11 @@ public class TransportRollupSearchAction extends TransportAction<SearchRequest,
                         String type = (String)agg.get(RollupField.AGG);
 
                         // If the cap is for a date_histo, and the query is a range, the timezones need to match
-                                if (type.equals(DateHistogramAggregationBuilder.NAME) && builder instanceof RangeQueryBuilder) {
+                        if (type.equals(DateHistogramAggregationBuilder.NAME) && timeZone != null) {
-                                    String timeZone = ((RangeQueryBuilder)builder).timeZone();
-
-                                    // Many range queries don't include the timezone because the default is UTC, but the query
-                                    // builder will return null so we need to set it here
-                                    if (timeZone == null) {
-                                        timeZone = DateTimeZone.UTC.toString();
-                                    }
                             boolean matchingTZ = ((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()))
                                 .equalsIgnoreCase(timeZone);
                             if (matchingTZ == false) {
-                                        incorrectTimeZones.add((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()));
+                                incompatibleTimeZones.add((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()));
                             }
                             return matchingTZ;
                         }
@@ -323,44 +354,20 @@ public class TransportRollupSearchAction extends TransportAction<SearchRequest,
             })
             .distinct()
             .collect(ArrayList::new, List::addAll, List::addAll);
-
+        if (rewrittenFieldNames.isEmpty()) {
-            if (rewrittenFieldName.isEmpty()) {
+            if (incompatibleTimeZones.isEmpty()) {
-                if (incorrectTimeZones.isEmpty()) {
+                throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builderName
-                    throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builder.getWriteableName()
                     + "] query is not available in selected rollup indices, cannot query.");
             } else {
-                    throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builder.getWriteableName()
+                throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builderName
                     + "] query was found in rollup indices, but requested timezone is not compatible. Options include: "
-                        + incorrectTimeZones);
+                    + incompatibleTimeZones);
             }
-            }
+        } else if (rewrittenFieldNames.size() > 1) {
-
-            if (rewrittenFieldName.size() > 1) {
             throw new IllegalArgumentException("Ambiguous field name resolution when mapping to rolled fields.  Field name [" +
-                        fieldName + "] was mapped to: [" + Strings.collectionToDelimitedString(rewrittenFieldName, ",") + "].");
+                fieldName + "] was mapped to: [" + Strings.collectionToDelimitedString(rewrittenFieldNames, ",") + "].");
-            }
-
-            //Note: instanceof here to make casting checks happier
-            if (builder instanceof RangeQueryBuilder) {
-                RangeQueryBuilder rewritten = new RangeQueryBuilder(rewrittenFieldName.get(0));
-                RangeQueryBuilder original = (RangeQueryBuilder)builder;
-                rewritten.from(original.from());
-                rewritten.to(original.to());
-                if (original.timeZone() != null) {
-                    rewritten.timeZone(original.timeZone());
-                }
-                rewritten.includeLower(original.includeLower());
-                rewritten.includeUpper(original.includeUpper());
-                return rewritten;
         } else {
-                return new TermQueryBuilder(rewrittenFieldName.get(0), ((TermQueryBuilder)builder).value());
+            return rewrittenFieldNames.get(0);
-            }
-
-        } else if (builder.getWriteableName().equals(MatchAllQueryBuilder.NAME)) {
-            // no-op
-            return builder;
-        } else {
-            throw new IllegalArgumentException("Unsupported Query in search request: [" + builder.getWriteableName() + "]");
         }
     }
 

x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/action/SearchActionTests.java

@@ -25,9 +25,11 @@ import org.elasticsearch.index.query.ConstantScoreQueryBuilder;
 import org.elasticsearch.index.query.DisMaxQueryBuilder;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.MatchPhraseQueryBuilder;
+import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.RangeQueryBuilder;
 import org.elasticsearch.index.query.TermQueryBuilder;
+import org.elasticsearch.index.query.TermsQueryBuilder;
 import org.elasticsearch.indices.IndicesModule;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.script.ScriptService;
@@ -61,6 +63,7 @@ import org.mockito.Mockito;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -153,7 +156,7 @@ public class SearchActionTests extends ESTestCase {
             "compatible. Options include: [UTC]"));
     }
 
-    public void testTerms() {
+    public void testTermQuery() {
         RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
         GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();
         group.setTerms(ConfigTestHelpers.getTerms().setFields(Collections.singletonList("foo")).build());
@@ -166,6 +169,23 @@ public class SearchActionTests extends ESTestCase {
         assertThat(((TermQueryBuilder)rewritten).fieldName(), equalTo("foo.terms.value"));
     }
 
+    public void testTermsQuery() {
+        RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
+        GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();
+        group.setTerms(ConfigTestHelpers.getTerms().setFields(Collections.singletonList("foo")).build());
+        job.setGroupConfig(group.build());
+        RollupJobCaps cap = new RollupJobCaps(job.build());
+        Set<RollupJobCaps> caps = new HashSet<>();
+        caps.add(cap);
+        QueryBuilder original = new TermsQueryBuilder("foo", Arrays.asList("bar", "baz"));
+        QueryBuilder rewritten =
+            TransportRollupSearchAction.rewriteQuery(original, caps);
+        assertThat(rewritten, instanceOf(TermsQueryBuilder.class));
+        assertNotSame(rewritten, original);
+        assertThat(((TermsQueryBuilder)rewritten).fieldName(), equalTo("foo.terms.value"));
+        assertThat(((TermsQueryBuilder)rewritten).values(),  equalTo(Arrays.asList("bar", "baz")));
+    }
+
     public void testCompounds() {
         RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
         GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();