Scripting: Remove unnecessary intermediate script compilation methods on QueryShardContext (#25093)

This commit removes wrapper methods on QueryShardContext used to compile
scripts. Instead, the script service is made accessible in the context,
and calls to compile can be made directly. This will ease transition to
each of those location becoming their own context, since they would no
longer be able to expect the same script class type.

11 files changed, 47 insertions, 72 deletions

diff --git a/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java b/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
index 68eb749e11..f13aa22f7d 100644
--- a/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
@@ -74,8 +74,9 @@ public abstract class InnerHitContextBuilder {
         }
         if (innerHitBuilder.getScriptFields() != null) {
             for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
-                SearchScript.LeafFactory searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
-                    SearchScript.CONTEXT);
+                QueryShardContext innerContext = innerHitsContext.getQueryShardContext();
+                SearchScript.Factory factory = innerContext.getScriptService().compile(field.script(), SearchScript.CONTEXT);
+                SearchScript.LeafFactory searchScript = factory.newFactory(field.script().getParams(), innerHitsContext.lookup());
                 innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
             }
diff --git a/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
index c95464776e..3fefb7141b 100644
--- a/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
+++ b/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
@@ -76,6 +76,11 @@ public class QueryRewriteContext {
         return mapperService;
     }
 
+    /** Return the script service to allow compiling scripts within queries. */
+    public ScriptService getScriptService() {
+        return scriptService;
+    }
+
     /** Return the current {@link IndexReader}, or {@code null} if no index reader is available, for
      *  instance if we are on the coordinating node or if this rewrite context is used to index
      *  queries (percolation). */
diff --git a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java b/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
index e82eb07199..2ce5abd213 100644
--- a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
+++ b/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
@@ -325,51 +325,9 @@ public class QueryShardContext extends QueryRewriteContext {
         return indexSettings.getIndex();
     }
 
-    /**
-     * Compiles (or retrieves from cache) and binds the parameters to the
-     * provided script
-     */
-    public final SearchScript.LeafFactory getSearchScript(Script script, ScriptContext<SearchScript.Factory> context) {
-        failIfFrozen();
-        SearchScript.Factory factory = scriptService.compile(script, context);
-        return factory.newFactory(script.getParams(), lookup());
-    }
-    /**
-     * Returns a lazily created {@link SearchScript} that is compiled immediately but can be pulled later once all
-     * parameters are available.
-     */
-    public final Function<Map<String, Object>, SearchScript.LeafFactory> getLazySearchScript(
-        Script script, ScriptContext<SearchScript.Factory> context) {
-        // TODO: this "lazy" binding can be removed once scripted metric aggs have their own contexts, which take _agg/_aggs as a parameter
-        failIfFrozen();
-        SearchScript.Factory factory = scriptService.compile(script, context);
-        return (p) -> factory.newFactory(p, lookup());
-    }
-
-    /**
-     * Compiles (or retrieves from cache) and binds the parameters to the
-     * provided script
-     */
-    public final ExecutableScript getExecutableScript(Script script, ScriptContext<ExecutableScript.Factory> context) {
-        failIfFrozen();
-        ExecutableScript.Factory factory = scriptService.compile(script, context);
-        return factory.newInstance(script.getParams());
-    }
-
-    /**
-     * Returns a lazily created {@link ExecutableScript} that is compiled immediately but can be pulled later once all
-     * parameters are available.
-     */
-    public final Function<Map<String, Object>, ExecutableScript> getLazyExecutableScript(
-        Script script, ScriptContext<ExecutableScript.Factory> context) {
-        // TODO: this "lazy" binding can be removed once scripted metric aggs have their own contexts, which take _agg/_aggs as a parameter
-        failIfFrozen();
-        ExecutableScript.Factory factory = scriptService.compile(script, context);
-        return factory::newInstance;
-    }
-
     /** Return the script service to allow compiling scripts. */
     public final ScriptService getScriptService() {
+        failIfFrozen();
         return scriptService;
     }
 
diff --git a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
index c86ff73d0f..0d608fd5f1 100644
--- a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
@@ -129,7 +129,9 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
-        return new ScriptQuery(script, context.getSearchScript(script, SearchScript.CONTEXT));
+        SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
+        SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
+        return new ScriptQuery(script, searchScript);
     }
 
     static class ScriptQuery extends Query {
diff --git a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java b/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
index 4c9ac0f452..7c66a7de99 100644
--- a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
@@ -94,7 +94,8 @@ public class ScriptScoreFunctionBuilder extends ScoreFunctionBuilder<ScriptScore
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
         try {
-            SearchScript.LeafFactory searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
+            SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
+            SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
             return new ScriptScoreFunction(script, searchScript);
         } catch (Exception e) {
             throw new QueryShardException(context, "script_score: the script could not be loaded", e);
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java
index e1fb72288b..99f3193b56 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java
@@ -27,6 +27,7 @@ import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.QueryParseContext;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.Script;
@@ -97,8 +98,9 @@ public class ScriptHeuristic extends SignificanceHeuristic {
 
     @Override
     public SignificanceHeuristic rewrite(SearchContext context) {
-        return new ExecutableScriptHeuristic(script,
-            context.getQueryShardContext().getExecutableScript(script, ExecutableScript.AGGS_CONTEXT));
+        QueryShardContext shardContext = context.getQueryShardContext();
+        ExecutableScript.Factory compiledScript = shardContext.getScriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
+        return new ExecutableScriptHeuristic(script, compiledScript.newInstance(script.getParams()));
     }
 
 
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
index 5b9e7e1dbd..4753884e43 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
@@ -184,22 +184,21 @@ public class ScriptedMetricAggregationBuilder extends AbstractAggregationBuilder
             Builder subfactoriesBuilder) throws IOException {
 
         QueryShardContext queryShardContext = context.getQueryShardContext();
-        Function<Map<String, Object>, ExecutableScript> executableInitScript;
+        ExecutableScript.Factory executableInitScript;
         if (initScript != null) {
-            executableInitScript = queryShardContext.getLazyExecutableScript(initScript, ExecutableScript.AGGS_CONTEXT);
+            executableInitScript = queryShardContext.getScriptService().compile(initScript, ExecutableScript.AGGS_CONTEXT);
         } else {
-            executableInitScript = (p) -> null;
+            executableInitScript = p -> null;
         }
-        Function<Map<String, Object>, SearchScript.LeafFactory> searchMapScript =
-            queryShardContext.getLazySearchScript(mapScript, SearchScript.AGGS_CONTEXT);
-        Function<Map<String, Object>, ExecutableScript> executableCombineScript;
+        SearchScript.Factory searchMapScript = queryShardContext.getScriptService().compile(mapScript, SearchScript.AGGS_CONTEXT);
+        ExecutableScript.Factory executableCombineScript;
         if (combineScript != null) {
-            executableCombineScript = queryShardContext.getLazyExecutableScript(combineScript, ExecutableScript.AGGS_CONTEXT);
+            executableCombineScript =queryShardContext.getScriptService().compile(combineScript, ExecutableScript.AGGS_CONTEXT);
         } else {
-            executableCombineScript = (p) -> null;
+            executableCombineScript = p -> null;
         }
         return new ScriptedMetricAggregatorFactory(name, searchMapScript, executableInitScript, executableCombineScript, reduceScript,
-                params, context, parent, subfactoriesBuilder, metaData);
+                params, queryShardContext.lookup(), context, parent, subfactoriesBuilder, metaData);
     }
 
 
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
index 6e75cba7b7..2d9e02d08c 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
@@ -28,6 +28,7 @@ import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -38,21 +39,23 @@ import java.util.function.Function;
 
 public class ScriptedMetricAggregatorFactory extends AggregatorFactory<ScriptedMetricAggregatorFactory> {
 
-    private final Function<Map<String, Object>, SearchScript.LeafFactory> mapScript;
-    private final Function<Map<String, Object>, ExecutableScript> combineScript;
+    private final SearchScript.Factory mapScript;
+    private final ExecutableScript.Factory combineScript;
     private final Script reduceScript;
     private final Map<String, Object> params;
-    private final Function<Map<String, Object>, ExecutableScript> initScript;
+    private final SearchLookup lookup;
+    private final ExecutableScript.Factory initScript;
 
-    public ScriptedMetricAggregatorFactory(String name, Function<Map<String, Object>, SearchScript.LeafFactory> mapScript,
-            Function<Map<String, Object>, ExecutableScript> initScript, Function<Map<String, Object>, ExecutableScript> combineScript,
-            Script reduceScript, Map<String, Object> params, SearchContext context, AggregatorFactory<?> parent,
-            AggregatorFactories.Builder subFactories, Map<String, Object> metaData) throws IOException {
+    public ScriptedMetricAggregatorFactory(String name, SearchScript.Factory mapScript, ExecutableScript.Factory initScript,
+                                           ExecutableScript.Factory combineScript, Script reduceScript, Map<String, Object> params,
+                                           SearchLookup lookup, SearchContext context, AggregatorFactory<?> parent,
+                                           AggregatorFactories.Builder subFactories, Map<String, Object> metaData) throws IOException {
         super(name, context, parent, subFactories, metaData);
         this.mapScript = mapScript;
         this.initScript = initScript;
         this.combineScript = combineScript;
         this.reduceScript = reduceScript;
+        this.lookup = lookup;
         this.params = params;
     }
 
@@ -70,9 +73,9 @@ public class ScriptedMetricAggregatorFactory extends AggregatorFactory<ScriptedM
             params.put("_agg", new HashMap<String, Object>());
         }
 
-        final ExecutableScript initScript = this.initScript.apply(params);
-        final SearchScript.LeafFactory mapScript = this.mapScript.apply(params);
-        final ExecutableScript combineScript = this.combineScript.apply(params);
+        final ExecutableScript initScript = this.initScript.newInstance(params);
+        final SearchScript.LeafFactory mapScript = this.mapScript.newFactory(params, lookup);
+        final ExecutableScript combineScript = this.combineScript.newInstance(params);
 
         final Script reduceScript = deepCopyScript(this.reduceScript, context);
         if (initScript != null) {
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
index 31cd1681f8..2739427cfe 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
@@ -27,6 +27,7 @@ import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.QueryParseContext;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
@@ -533,8 +534,9 @@ public class TopHitsAggregationBuilder extends AbstractAggregationBuilder<TopHit
         List<ScriptFieldsContext.ScriptField> fields = new ArrayList<>();
         if (scriptFields != null) {
             for (ScriptField field : scriptFields) {
-                SearchScript.LeafFactory searchScript = context.getQueryShardContext().getSearchScript(field.script(),
-                    SearchScript.CONTEXT);
+                QueryShardContext shardContext = context.getQueryShardContext();
+                SearchScript.Factory factory = shardContext.getScriptService().compile(field.script(), SearchScript.CONTEXT);
+                SearchScript.LeafFactory searchScript = factory.newFactory(field.script().getParams(), shardContext.lookup());
                 fields.add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
             }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
index d7b6272308..d8c2167c16 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
@@ -120,7 +120,8 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
         if (script == null) {
             return null;
         } else {
-            return context.getSearchScript(script, SearchScript.AGGS_CONTEXT);
+            SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.AGGS_CONTEXT);
+            return factory.newFactory(script.getParams(), context.lookup());
         }
     }
 
diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
index 18b15bfec2..749d2d13ba 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
@@ -240,7 +240,8 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
 
     @Override
     public SortFieldAndFormat build(QueryShardContext context) throws IOException {
-        final SearchScript.LeafFactory searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
+        final SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
+        final SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
 
         MultiValueMode valueMode = null;
         if (sortMode != null) {