Move query builder caching check to dedicated tests (#43238)

Currently `AbstractQueryTestCase#testToQuery` checks the search context cachable
flag. This is a bit fragile due to the high randomization of query builders
performed by this general test. Also we might only rarely check the
"interesting" cases because they rarely get generated when fully randomizing the
query builder.

This change moved the general checks out ot #testToQuery and instead adds
dedicated cache tests for those query builders that exhibit something other than
the default behaviour.

Closes #43200

modules/percolator/src/test/java/org/elasticsearch/percolator/PercolateQueryBuilderTests.java

@@ -296,9 +296,17 @@ private static BytesReference randomSource(Set<String> usedFields) {
         }
     }
 
+    /**
+     * Test that this query is never cacheable
+     */
     @Override
-    protected boolean isCacheable(PercolateQueryBuilder queryBuilder) {
-        return false;
+    public void testCacheability() throws IOException {
+        PercolateQueryBuilder queryBuilder = createTestQueryBuilder();
+        QueryShardContext context = createShardContext();
+        assert context.isCacheable();
+        QueryBuilder rewritten = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewritten.toQuery(context));
+        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreQueryBuilder.java

@@ -23,17 +23,17 @@
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.lucene.search.function.ScriptScoreFunction;
+import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
-import org.elasticsearch.index.query.QueryBuilder;
-import org.elasticsearch.script.Script;
-import org.elasticsearch.common.lucene.search.function.ScriptScoreFunction;
-import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
 import org.elasticsearch.index.query.InnerHitContextBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.script.Script;
 
 import java.io.IOException;
 import java.util.Map;

server/src/test/java/org/elasticsearch/index/query/MoreLikeThisQueryBuilderTests.java

@@ -363,9 +363,34 @@ public void testItemFromXContent() throws IOException {
         assertEquals(expectedItem, newItem);
     }
 
+    /**
+     * Check that this query is generally not cacheable, except when we fetch 0 items
+     */
     @Override
-    protected boolean isCacheable(MoreLikeThisQueryBuilder queryBuilder) {
-        return queryBuilder.likeItems().length == 0; // items are always fetched
+    public void testCacheability() throws IOException {
+        MoreLikeThisQueryBuilder queryBuilder = createTestQueryBuilder();
+        boolean isCacheable = queryBuilder.likeItems().length == 0; // items are always fetched
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertEquals("query should " + (isCacheable ? "" : "not") + " be cacheable: " + queryBuilder.toString(), isCacheable,
+                context.isCacheable());
+
+        // specifically trigger case where query is cacheable
+        queryBuilder = new MoreLikeThisQueryBuilder(randomStringFields(), new String[] {"some text"}, null);
+        context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
+
+        // specifically trigger case where query is not cacheable
+        queryBuilder = new MoreLikeThisQueryBuilder(randomStringFields(), null, new Item[] { new Item("foo", "1") });
+        context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
     public void testFromJson() throws IOException {
@@ -405,8 +430,6 @@ public void testFromJson() throws IOException {
     protected QueryBuilder parseQuery(XContentParser parser) throws IOException {
         QueryBuilder query = super.parseQuery(parser);
         assertThat(query, instanceOf(MoreLikeThisQueryBuilder.class));
-
-        MoreLikeThisQueryBuilder mltQuery = (MoreLikeThisQueryBuilder) query;
         return query;
     }
 

server/src/test/java/org/elasticsearch/index/query/RangeQueryBuilderTests.java

@@ -601,4 +601,25 @@ public void testTypeField() throws IOException {
         builder.doToQuery(createShardContext());
         assertWarnings(QueryShardContext.TYPES_DEPRECATION_MESSAGE);
     }
+
+    /**
+     * Range queries should generally be cacheable, at least the ones we create randomly.
+     * This test makes sure we also test the non-cacheable cases regularly.
+     */
+    @Override
+    public void testCacheability() throws IOException {
+        RangeQueryBuilder queryBuilder = createTestQueryBuilder();
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
+
+        // queries on date fields using "now" should not be cached
+        queryBuilder = new RangeQueryBuilder(randomFrom(DATE_FIELD_NAME, DATE_RANGE_FIELD_NAME, DATE_ALIAS_FIELD_NAME));
+        queryBuilder.to("now");
+        context = createShardContext();
+        rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
+    }
 }

server/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java

@@ -116,8 +116,15 @@ protected Set<String> getObjectsHoldingArbitraryContent() {
         return Collections.singleton(Script.PARAMS_PARSE_FIELD.getPreferredName());
     }
 
+    /**
+     * Check that this query is generally not cacheable
+     */
     @Override
-    protected boolean isCacheable(ScriptQueryBuilder queryBuilder) {
-        return false;
+    public void testCacheability() throws IOException {
+        ScriptQueryBuilder queryBuilder = createTestQueryBuilder();
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 }

server/src/test/java/org/elasticsearch/index/query/ScriptScoreQueryBuilderTests.java

@@ -88,8 +88,15 @@ public void testIllegalArguments() {
         );
     }
 
+    /**
+     * Check that this query is generally not cacheable
+     */
     @Override
-    protected boolean isCacheable(ScriptScoreQueryBuilder queryBuilder) {
-        return false;
+    public void testCacheability() throws IOException {
+        ScriptScoreQueryBuilder queryBuilder = createTestQueryBuilder();
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 }

server/src/test/java/org/elasticsearch/index/query/TermsQueryBuilderTests.java

@@ -281,13 +281,6 @@ public void testGeo() throws Exception {
                 e.getMessage());
     }
 
-    @Override
-    protected boolean isCacheable(TermsQueryBuilder queryBuilder) {
-        // even though we use a terms lookup here we do this during rewrite and that means we are cacheable on toQuery
-        // that's why we return true here all the time
-        return super.isCacheable(queryBuilder);
-    }
-
     public void testSerializationFailsUnlessFetched() throws IOException {
         QueryBuilder builder = new TermsQueryBuilder(STRING_FIELD_NAME, randomTermsLookup());
         QueryBuilder termsQueryBuilder = Rewriteable.rewrite(builder, createShardContext());

server/src/test/java/org/elasticsearch/index/query/TermsSetQueryBuilderTests.java

@@ -110,10 +110,42 @@ protected void doAssertLuceneQuery(TermsSetQueryBuilder queryBuilder, Query quer
         }
     }
 
+    /**
+     * Check that this query is generally not cacheable and explicitly testing the two conditions when it is not as well
+     */
     @Override
-    protected boolean isCacheable(TermsSetQueryBuilder queryBuilder) {
-        return queryBuilder.getMinimumShouldMatchField() != null ||
+    public void testCacheability() throws IOException {
+        TermsSetQueryBuilder queryBuilder = createTestQueryBuilder();
+        boolean isCacheable = queryBuilder.getMinimumShouldMatchField() != null ||
                 (queryBuilder.getMinimumShouldMatchScript() != null && queryBuilder.getValues().isEmpty());
+        QueryShardContext context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(queryBuilder.doToQuery(context));
+        assertEquals("query should " + (isCacheable ? "" : "not") + " be cacheable: " + queryBuilder.toString(), isCacheable,
+                context.isCacheable());
+
+        // specifically trigger the two cases where query is cacheable
+        queryBuilder = new TermsSetQueryBuilder(STRING_FIELD_NAME, Collections.singletonList("foo"));
+        queryBuilder.setMinimumShouldMatchField("m_s_m");
+        context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(queryBuilder.doToQuery(context));
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
+
+        queryBuilder = new TermsSetQueryBuilder(STRING_FIELD_NAME, Collections.emptyList());
+        queryBuilder.setMinimumShouldMatchScript(new Script(ScriptType.INLINE, MockScriptEngine.NAME, "_script", emptyMap()));
+        context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(queryBuilder.doToQuery(context));
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
+
+        // also test one case where query is not cacheable
+        queryBuilder = new TermsSetQueryBuilder(STRING_FIELD_NAME, Collections.singletonList("foo"));
+        queryBuilder.setMinimumShouldMatchScript(new Script(ScriptType.INLINE, MockScriptEngine.NAME, "_script", emptyMap()));
+        context = createShardContext();
+        rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(queryBuilder.doToQuery(context));
+        assertFalse("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
     @Override

server/src/test/java/org/elasticsearch/index/query/TypeQueryBuilderTests.java

@@ -83,4 +83,10 @@ public void testMustRewrite() throws IOException {
         super.testMustRewrite();
         assertWarnings(TypeQueryBuilder.TYPES_DEPRECATION_MESSAGE);
     }
+
+    @Override
+    public void testCacheability() throws IOException {
+        super.testCacheability();
+        assertWarnings(TypeQueryBuilder.TYPES_DEPRECATION_MESSAGE);
+    }
 }

server/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.index.query.functionscore;
 
 import com.fasterxml.jackson.core.JsonParseException;
+
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
@@ -40,6 +41,7 @@
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.RandomQueryBuilder;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.index.query.WrapperQueryBuilder;
@@ -676,7 +678,7 @@ public void testRewriteWithFunction() throws IOException {
      */
     public void testSingleScriptFunction() throws IOException {
         QueryBuilder queryBuilder = RandomQueryBuilder.createQuery(random());
-        ScoreFunctionBuilder functionBuilder = new ScriptScoreFunctionBuilder(
+        ScoreFunctionBuilder<ScriptScoreFunctionBuilder> functionBuilder = new ScriptScoreFunctionBuilder(
             new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));
 
         FunctionScoreQueryBuilder builder = functionScoreQuery(queryBuilder, functionBuilder);
@@ -796,8 +798,38 @@ public List<ScoreFunctionSpec<?>> getScoreFunctions() {
         }
     }
 
+    /**
+     * Check that this query is generally cacheable except for builders using {@link ScriptScoreFunctionBuilder} or
+     * {@link RandomScoreFunctionBuilder} without a seed
+     */
     @Override
-    protected boolean isCacheable(FunctionScoreQueryBuilder queryBuilder) {
+    public void testCacheability() throws IOException {
+        FunctionScoreQueryBuilder queryBuilder = createTestQueryBuilder();
+        boolean isCacheable = isCacheable(queryBuilder);
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertEquals("query should " + (isCacheable ? "" : "not") + " be cacheable: " + queryBuilder.toString(), isCacheable,
+                context.isCacheable());
+
+        // check the two non-cacheable cases explicitly
+        ScoreFunctionBuilder<?> scriptScoreFunction = new ScriptScoreFunctionBuilder(
+                new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));
+        RandomScoreFunctionBuilder randomScoreFunctionBuilder = new RandomScoreFunctionBuilderWithFixedSeed();
+
+        for (ScoreFunctionBuilder<?> scoreFunction : List.of(scriptScoreFunction, randomScoreFunctionBuilder)) {
+            FilterFunctionBuilder[] functions = new FilterFunctionBuilder[] {
+                    new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), scoreFunction) };
+            queryBuilder = new FunctionScoreQueryBuilder(functions);
+
+            context = createShardContext();
+            rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+            assertNotNull(rewriteQuery.toQuery(context));
+            assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
+        }
+    }
+
+    private boolean isCacheable(FunctionScoreQueryBuilder queryBuilder) {
         FilterFunctionBuilder[] filterFunctionBuilders = queryBuilder.filterFunctionBuilders();
         for (FilterFunctionBuilder builder : filterFunctionBuilders) {
             if (builder.getScoreFunction() instanceof ScriptScoreFunctionBuilder) {

test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java

@@ -428,13 +428,6 @@ public void testToQuery() throws IOException {
              * we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
             QueryBuilder rewritten = rewriteQuery(firstQuery, new QueryShardContext(context));
             Query firstLuceneQuery = rewritten.toQuery(context);
-            if (isCacheable(firstQuery)) {
-                assertTrue("query was marked as not cacheable in the context but this test indicates it should be cacheable: "
-                        + firstQuery.toString(), context.isCacheable());
-            } else {
-                assertFalse("query was marked as cacheable in the context but this test indicates it should not be cacheable: "
-                        + firstQuery.toString(), context.isCacheable());
-            }
             assertNotNull("toQuery should not return null", firstLuceneQuery);
             assertLuceneQuery(firstQuery, firstLuceneQuery, searchContext);
             //remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well
@@ -478,10 +471,6 @@ protected QueryBuilder rewriteQuery(QB queryBuilder, QueryRewriteContext rewrite
         return rewritten;
     }
 
-    protected boolean isCacheable(QB queryBuilder) {
-        return true;
-    }
-
     /**
      * Few queries allow you to set the boost on the Java API, although the corresponding parser
      * doesn't parse it as it isn't supported. This method allows to disable boost related tests for those queries.
@@ -806,4 +795,17 @@ protected QueryBuilder rewriteAndFetch(QueryBuilder builder, QueryRewriteContext
     public boolean isTextField(String fieldName) {
         return fieldName.equals(STRING_FIELD_NAME) || fieldName.equals(STRING_ALIAS_FIELD_NAME);
     }
+
+    /**
+     * Check that a query is generally cacheable. Tests for query builders that are not always cacheable
+     * should overwrite this method and make sure the different cases are always tested
+     */
+    public void testCacheability() throws IOException {
+        QB queryBuilder = createTestQueryBuilder();
+        QueryShardContext context = createShardContext();
+        QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
+    }
+
 }