From 48cd016f909bbeca7b0d0a4926e36a333d021869 Mon Sep 17 00:00:00 2001
From: Daniel Worley <danworley@gmail.com>
Date: Tue, 27 Apr 2021 16:56:33 -0400
Subject: [PATCH] Es 7 12 (#364)

* wip for 7.11 deprecations

* Fix build errors for 7.11

* add public visibility to test templates

* keeping threadpool over executor for now

* remove all the checks except keep nebula disabled

* Restore override annotation in TransportAddFeatureToSet

* read() was removed

* Bring back checkstyle and fix some violations

* Clean up unused vars

* 7.12 changes

Co-authored-by: Nathan Day <nathancday@gmail.com>
---
 gradle.properties                             |  4 ++--
 .../o19s/es/explore/ExplorerQueryBuilder.java |  4 ++--
 .../java/com/o19s/es/ltr/LtrQueryContext.java | 10 ++++-----
 .../store/PrecompiledTemplateFeature.java     |  7 +++---
 .../es/ltr/feature/store/ScriptFeature.java   | 14 ++++++------
 .../es/ltr/feature/store/StoredFeature.java   |  5 +++--
 .../o19s/es/ltr/query/LtrQueryBuilder.java    |  4 ++--
 .../es/ltr/query/StoredLtrQueryBuilder.java   |  6 ++---
 .../ltr/query/ValidatingLtrQueryBuilder.java  |  8 +++----
 .../es/termstat/TermStatQueryBuilder.java     |  8 +++----
 .../es/explore/ExplorerQueryBuilderTests.java |  6 ++---
 .../es/ltr/query/LtrQueryBuilderTests.java    | 14 ++++++------
 .../ltr/query/StoredLtrQueryBuilderTests.java | 22 +++++++++----------
 .../query/ValidatingLtrQueryBuilderTests.java |  5 +++--
 .../termstat/TermStatQueryBuilderTests.java   |  4 ++--
 15 files changed, 62 insertions(+), 59 deletions(-)

diff --git a/gradle.properties b/gradle.properties
index 10717d12..c8416f9c 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -1,5 +1,5 @@
 ltrVersion = 1.5.4
-elasticsearchVersion = 7.11.1
-luceneVersion = 8.7.0
+elasticsearchVersion = 7.12.0
+luceneVersion = 8.8.0
 ow2Version = 8.0.1
 antlrVersion = 4.5.1-1
diff --git a/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java b/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java
index 8e4fb04c..86ded2c7 100644
--- a/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java
+++ b/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java
@@ -27,7 +27,7 @@
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.Rewriteable;
 
 import java.io.IOException;
@@ -98,7 +98,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     }
 
     @Override
-    protected Query doToQuery(QueryShardContext context) throws IOException {
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return new ExplorerQuery(query.toQuery(context), type);
     }
 
diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryContext.java b/src/main/java/com/o19s/es/ltr/LtrQueryContext.java
index 7d10f2d7..2912a02f 100644
--- a/src/main/java/com/o19s/es/ltr/LtrQueryContext.java
+++ b/src/main/java/com/o19s/es/ltr/LtrQueryContext.java
@@ -16,7 +16,7 @@
 
 package com.o19s.es.ltr;
 
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 
 import java.util.Collections;
 import java.util.Set;
@@ -25,19 +25,19 @@
  * LTR queryShardContext used to track information needed for building lucene queries
  */
 public class LtrQueryContext {
-    private final QueryShardContext queryShardContext;
+    private final SearchExecutionContext queryShardContext;
     private final Set<String> activeFeatures;
 
-    public LtrQueryContext(QueryShardContext context) {
+    public LtrQueryContext(SearchExecutionContext context) {
         this(context, Collections.emptySet());
     }
 
-    public LtrQueryContext(QueryShardContext context, Set<String> activeFeatures) {
+    public LtrQueryContext(SearchExecutionContext context, Set<String> activeFeatures) {
         this.queryShardContext = context;
         this.activeFeatures = activeFeatures;
     }
 
-    public QueryShardContext getQueryShardContext() {
+    public SearchExecutionContext getSearchExecutionContext() {
         return queryShardContext;
     }
 
diff --git a/src/main/java/com/o19s/es/ltr/feature/store/PrecompiledTemplateFeature.java b/src/main/java/com/o19s/es/ltr/feature/store/PrecompiledTemplateFeature.java
index 14d7a3f0..4d887275 100644
--- a/src/main/java/com/o19s/es/ltr/feature/store/PrecompiledTemplateFeature.java
+++ b/src/main/java/com/o19s/es/ltr/feature/store/PrecompiledTemplateFeature.java
@@ -93,14 +93,15 @@ public Query doToQuery(LtrQueryContext context, FeatureSet set, Map<String, Obje
         String query = MustacheUtils.execute(template, params);
         try {
             XContentParser parser = XContentFactory.xContent(query)
-                    .createParser(context.getQueryShardContext().getXContentRegistry(),
+                    .createParser(context.getSearchExecutionContext().getXContentRegistry(),
                             LoggingDeprecationHandler.INSTANCE, query);
             QueryBuilder queryBuilder = parseInnerQueryBuilder(parser);
             // XXX: QueryShardContext extends QueryRewriteContext (for now)
-            return Rewriteable.rewrite(queryBuilder, context.getQueryShardContext()).toQuery(context.getQueryShardContext());
+            return Rewriteable.rewrite(queryBuilder, context.getSearchExecutionContext()).toQuery(context.getSearchExecutionContext());
         } catch (IOException | ParsingException | IllegalArgumentException e) {
             // wrap common exceptions as well so we can attach the feature's name to the stack
-            throw new QueryShardException(context.getQueryShardContext(), "Cannot create query while parsing feature [" + name + "]", e);
+            throw new QueryShardException(context.getSearchExecutionContext(),
+                    "Cannot create query while parsing feature [" + name + "]", e);
         }
     }
 
diff --git a/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java b/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java
index 2c813056..bcffac28 100644
--- a/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java
+++ b/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java
@@ -173,10 +173,10 @@ public Query doToQuery(LtrQueryContext context, FeatureSet featureSet, Map<Strin
             Analyzer analyzer = null;
             for(String field : fields) {
                 if (analyzerName == null) {
-                    final MappedFieldType fieldType = context.getQueryShardContext().getFieldType(field);
+                    final MappedFieldType fieldType = context.getSearchExecutionContext().getFieldType(field);
                     analyzer = fieldType.getTextSearchInfo().getSearchAnalyzer();
                 } else {
-                    analyzer = context.getQueryShardContext().getIndexAnalyzers().get(analyzerName);
+                    analyzer = context.getSearchExecutionContext().getIndexAnalyzers().get(analyzerName);
                 }
 
                 if (analyzer == null) {
@@ -211,12 +211,12 @@ public Query doToQuery(LtrQueryContext context, FeatureSet featureSet, Map<Strin
         nparams.put(EXTRA_LOGGING, extraLoggingSupplier);
         Script script = new Script(this.script.getType(), this.script.getLang(),
             this.script.getIdOrCode(), this.script.getOptions(), nparams);
-        ScoreScript.Factory factoryFactory  = context.getQueryShardContext().compile(script, ScoreScript.CONTEXT);
-        ScoreScript.LeafFactory leafFactory = factoryFactory.newFactory(nparams, context.getQueryShardContext().lookup());
+        ScoreScript.Factory factoryFactory  = context.getSearchExecutionContext().compile(script, ScoreScript.CONTEXT);
+        ScoreScript.LeafFactory leafFactory = factoryFactory.newFactory(nparams, context.getSearchExecutionContext().lookup());
         ScriptScoreFunction function = new ScriptScoreFunction(script, leafFactory,
-                context.getQueryShardContext().index().getName(),
-                context.getQueryShardContext().getShardId(),
-                context.getQueryShardContext().indexVersionCreated());
+                context.getSearchExecutionContext().index().getName(),
+                context.getSearchExecutionContext().getShardId(),
+                context.getSearchExecutionContext().indexVersionCreated());
         return new LtrScript(function, supplier, extraLoggingSupplier, termstatSupplier, terms);
     }
 
diff --git a/src/main/java/com/o19s/es/ltr/feature/store/StoredFeature.java b/src/main/java/com/o19s/es/ltr/feature/store/StoredFeature.java
index abf8bc56..b1cdbe2b 100644
--- a/src/main/java/com/o19s/es/ltr/feature/store/StoredFeature.java
+++ b/src/main/java/com/o19s/es/ltr/feature/store/StoredFeature.java
@@ -210,10 +210,11 @@ public Query doToQuery(LtrQueryContext context, FeatureSet set, Map<String, Obje
         // TODO: verify that this actually works, it does not feel right
         ScriptQueryBuilder builder = new ScriptQueryBuilder(new Script(ScriptType.INLINE, templateLanguage, template, params));
         try {
-            return builder.toQuery(context.getQueryShardContext());
+            return builder.toQuery(context.getSearchExecutionContext());
         } catch (IOException | ParsingException | IllegalArgumentException e) {
             // wrap common exceptions as well so we can attach the feature's name to the stack
-            throw new QueryShardException(context.getQueryShardContext(), "Cannot create query while parsing feature [" + name + "]", e);
+            throw new QueryShardException(context.getSearchExecutionContext(),
+                    "Cannot create query while parsing feature [" + name + "]", e);
         }
     }
 
diff --git a/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java
index 8a280072..18d0d130 100644
--- a/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java
+++ b/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java
@@ -35,7 +35,7 @@
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.script.Script;
 
@@ -122,7 +122,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     }
 
     @Override
-    protected Query doToQuery(QueryShardContext context) throws IOException {
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
         List<PrebuiltFeature> features = new ArrayList<>(_features.size());
         for (QueryBuilder builder : _features) {
             features.add(new PrebuiltFeature(builder.queryName(), builder.toQuery(context)));
diff --git a/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java
index 2d18239b..dc59dbf0 100644
--- a/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java
+++ b/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java
@@ -33,7 +33,7 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -66,7 +66,7 @@ public class StoredLtrQueryBuilder extends AbstractQueryBuilder<StoredLtrQueryBu
     }
 
     /**
-     * Injected context used to load a {@link FeatureStore} when running {@link #doToQuery(QueryShardContext)}
+     * Injected context used to load a {@link FeatureStore} when running {@link #doToQuery(SearchExecutionContext)}
      */
     private final transient FeatureStoreLoader storeLoader;
     private String modelName;
@@ -153,7 +153,7 @@ private static void validateActiveFeatures(FeatureSet features, LtrQueryContext
     }
 
     @Override
-    protected RankerQuery doToQuery(QueryShardContext context) throws IOException {
+    protected RankerQuery doToQuery(SearchExecutionContext context) throws IOException {
         String indexName = storeName != null ? IndexFeatureStore.indexName(storeName) : IndexFeatureStore.DEFAULT_STORE;
         FeatureStore store = storeLoader.load(indexName, context::getClient);
         LtrQueryContext ltrQueryContext = new LtrQueryContext(context,
diff --git a/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java
index 11ca3dcf..83323081 100644
--- a/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java
+++ b/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java
@@ -38,7 +38,7 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.QueryShardException;
 
 import java.io.IOException;
@@ -158,9 +158,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     }
 
     @Override
-    protected Query doToQuery(QueryShardContext queryShardContext) throws IOException {
+    protected Query doToQuery(SearchExecutionContext searchExecutionContext) throws IOException {
         //TODO: should we be passing activeFeatures here?
-        LtrQueryContext context = new LtrQueryContext(queryShardContext);
+        LtrQueryContext context = new LtrQueryContext(searchExecutionContext);
         if (StoredFeature.TYPE.equals(element.type())) {
             Feature feature = ((StoredFeature) element).optimize();
             if (feature instanceof PrecompiledExpressionFeature) {
@@ -178,7 +178,7 @@ protected Query doToQuery(QueryShardContext queryShardContext) throws IOExceptio
             CompiledLtrModel model = ((StoredLtrModel) element).compile(factory);
             return RankerQuery.build(model, context, validation.getParams());
         } else {
-            throw new QueryShardException(queryShardContext, "Unknown element type [" + element.type() + "]");
+            throw new QueryShardException(searchExecutionContext, "Unknown element type [" + element.type() + "]");
         }
     }
 
diff --git a/src/main/java/com/o19s/es/termstat/TermStatQueryBuilder.java b/src/main/java/com/o19s/es/termstat/TermStatQueryBuilder.java
index 0469992a..8f01f456 100644
--- a/src/main/java/com/o19s/es/termstat/TermStatQueryBuilder.java
+++ b/src/main/java/com/o19s/es/termstat/TermStatQueryBuilder.java
@@ -19,7 +19,7 @@
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -134,7 +134,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     }
 
     @Override
-    protected Query doToQuery(QueryShardContext context) throws IOException {
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
         Expression compiledExpression = (Expression) Scripting.compile(expr);
         AggrType aggrType = AggrType.valueOf(aggr.toUpperCase(Locale.getDefault()));
         AggrType posAggrType = AggrType.valueOf(pos_aggr.toUpperCase(Locale.getDefault()));
@@ -169,12 +169,12 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
         return new TermStatQuery(compiledExpression, aggrType, posAggrType, termSet);
     }
 
-    private Analyzer getAnalyzerForField(QueryShardContext context, String fieldName) {
+    private Analyzer getAnalyzerForField(SearchExecutionContext context, String fieldName) {
         MappedFieldType fieldType = context.getFieldType(fieldName);
         return fieldType.getTextSearchInfo().getSearchAnalyzer();
     }
 
-    private Analyzer getAnalyzerByName(QueryShardContext context, String analyzerName) {
+    private Analyzer getAnalyzerByName(SearchExecutionContext context, String analyzerName) {
         return context.getIndexAnalyzers().get(analyzerName);
     }
 
diff --git a/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java b/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java
index 594a68cb..65794abc 100644
--- a/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java
+++ b/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java
@@ -20,7 +20,7 @@
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.AbstractQueryTestCase;
@@ -66,7 +66,7 @@ public void testParse() throws Exception {
 
     @Override
     public void testMustRewrite() throws IOException {
-        QueryShardContext context = createShardContext();
+        SearchExecutionContext context = createSearchExecutionContext();
         context.setAllowUnmappedFields(true);
         ExplorerQueryBuilder queryBuilder = createTestQueryBuilder();
         queryBuilder.boost(AbstractQueryBuilder.DEFAULT_BOOST);
@@ -104,7 +104,7 @@ public void testMissingType() throws Exception {
     }
 
     @Override
-    protected void doAssertLuceneQuery(ExplorerQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
+    protected void doAssertLuceneQuery(ExplorerQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
         assertThat(query, instanceOf(ExplorerQuery.class));
     }
 }
diff --git a/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java
index 8eb60107..f23742ad 100644
--- a/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java
+++ b/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java
@@ -21,7 +21,7 @@
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.MatchQueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.WrapperQueryBuilder;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.Script;
@@ -125,7 +125,7 @@ public void testNamedFeatures() throws IOException {
                 "   } " +
                 "}";
         LtrQueryBuilder queryBuilder = (LtrQueryBuilder)parseQuery(ltrQuery);
-        QueryShardContext context = createShardContext();
+        SearchExecutionContext context = createSearchExecutionContext();
         RankerQuery query = (RankerQuery)queryBuilder.toQuery(context);
         assertEquals(query.getFeature(0).name(), "bar_query");
         assertEquals(query.getFeature(1).name(), "sham_query");
@@ -152,7 +152,7 @@ public void testUnnamedFeatures() throws IOException {
                 "   } " +
                 "}";
         LtrQueryBuilder queryBuilder = (LtrQueryBuilder)parseQuery(ltrQuery);
-        QueryShardContext context = createShardContext();
+        SearchExecutionContext context = createSearchExecutionContext();
         RankerQuery query = (RankerQuery)queryBuilder.toQuery(context);
         assertNull(query.getFeature(0).name());
         assertEquals(query.getFeature(1).name(), "");
@@ -167,9 +167,9 @@ protected boolean builderGeneratesCacheableQueries() {
     @Override
     public void testCacheability() throws IOException {
         LtrQueryBuilder queryBuilder = createTestQueryBuilder();
-        QueryShardContext context = createShardContext();
+        SearchExecutionContext context = createSearchExecutionContext();
         assert context.isCacheable();
-        QueryBuilder rewritten = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        QueryBuilder rewritten = rewriteQuery(queryBuilder, new SearchExecutionContext(context));
         assertNotNull(rewritten.toQuery(context));
         assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
@@ -215,7 +215,7 @@ public void testMustRewrite() throws IOException {
         }
 
         LtrQueryBuilder builder = new LtrQueryBuilder(script, features);
-        QueryBuilder rewritten = builder.rewrite(createShardContext());
+        QueryBuilder rewritten = builder.rewrite(createSearchExecutionContext());
         if (!mustRewrite && features.isEmpty()) {
             // if it's empty we rewrite to match all
             assertEquals(rewritten, new MatchAllQueryBuilder());
@@ -235,7 +235,7 @@ public void testMustRewrite() throws IOException {
     }
 
     @Override
-    protected void doAssertLuceneQuery(LtrQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
+    protected void doAssertLuceneQuery(LtrQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
         assertThat(query, instanceOf(RankerQuery.class));
     }
 
diff --git a/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java
index 4b937ea2..a40c151f 100644
--- a/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java
+++ b/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java
@@ -40,7 +40,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.query.functionscore.FieldValueFactorFunctionBuilder;
 import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
@@ -132,13 +132,13 @@ public void testMissingParams() {
         StoredLtrQueryBuilder builder = new StoredLtrQueryBuilder(LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store));
         builder.modelName("model1");
 
-        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createShardContext())).getMessage(),
+        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createSearchExecutionContext())).getMessage(),
                 equalTo("Missing required param(s): [query_string]"));
 
         Map<String, Object> params = new HashMap<>();
         params.put("query_string2", "a wonderful query");
         builder.params(params);
-        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createShardContext())).getMessage(),
+        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createSearchExecutionContext())).getMessage(),
                 equalTo("Missing required param(s): [query_string]"));
 
     }
@@ -147,7 +147,7 @@ public void testInvalidActiveFeatures() {
         StoredLtrQueryBuilder builder = new StoredLtrQueryBuilder(LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store));
         builder.modelName("model1");
         builder.activeFeatures(Collections.singletonList("non_existent_feature"));
-        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createShardContext())).getMessage(),
+        assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createSearchExecutionContext())).getMessage(),
                 equalTo("Feature: [non_existent_feature] provided in active_features does not exist"));
     }
 
@@ -184,14 +184,14 @@ private void assertQueryClass(Class<?> clazz, boolean setActiveFeature) throws I
             builder.activeFeatures(Arrays.asList("match1", "match2"));
         }
 
-        RankerQuery rankerQuery = builder.doToQuery(createShardContext());
+        RankerQuery rankerQuery = builder.doToQuery(createSearchExecutionContext());
         List<Query> queries = rankerQuery.stream().collect(Collectors.toList());
         assertEquals(clazz, queries.get(2).getClass());
     }
 
     @Override
     protected void doAssertLuceneQuery(StoredLtrQueryBuilder queryBuilder,
-                                       Query query, QueryShardContext context) throws IOException {
+                                       Query query, SearchExecutionContext context) throws IOException {
         assertThat(query, instanceOf(RankerQuery.class));
         RankerQuery rquery = (RankerQuery) query;
         Iterator<Query> ite = rquery.stream().iterator();
@@ -209,7 +209,7 @@ protected void doAssertLuceneQuery(StoredLtrQueryBuilder queryBuilder,
         assertTrue(ite.hasNext());
         Query featureQuery = ite.next();
         QueryBuilder builder = new MatchQueryBuilder("field1", queryBuilder.params().get("query_string"));
-        QueryShardContext qcontext = createShardContext();
+        SearchExecutionContext qcontext = createSearchExecutionContext();
 
         Query expected = Rewriteable.rewrite(builder, qcontext).toQuery(qcontext);
         assertEquals(expected, featureQuery);
@@ -217,7 +217,7 @@ protected void doAssertLuceneQuery(StoredLtrQueryBuilder queryBuilder,
         assertTrue(ite.hasNext());
         featureQuery = ite.next();
         builder = new MatchQueryBuilder("field2", queryBuilder.params().get("query_string"));
-        qcontext = createShardContext();
+        qcontext = createSearchExecutionContext();
         expected = Rewriteable.rewrite(builder, qcontext).toQuery(qcontext);
         assertEquals(expected, featureQuery);
         assertTrue(ite.hasNext());
@@ -227,7 +227,7 @@ protected void doAssertLuceneQuery(StoredLtrQueryBuilder queryBuilder,
                 .factor(1.2F)
                 .modifier(FieldValueFactorFunction.Modifier.LN2P)
                 .missing(0F));
-        qcontext = createShardContext();
+        qcontext = createSearchExecutionContext();
         expected = Rewriteable.rewrite(builder, qcontext).toQuery(qcontext);
         assertEquals(expected, featureQuery);
 
@@ -237,9 +237,9 @@ protected void doAssertLuceneQuery(StoredLtrQueryBuilder queryBuilder,
     @Override
     public void testCacheability() throws IOException {
         StoredLtrQueryBuilder queryBuilder = createTestQueryBuilder();
-        QueryShardContext context = createShardContext();
+        SearchExecutionContext context = createSearchExecutionContext();
         assert context.isCacheable();
-        QueryBuilder rewritten = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        QueryBuilder rewritten = rewriteQuery(queryBuilder, new SearchExecutionContext(context));
         assertNotNull(rewritten.toQuery(context));
         assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
diff --git a/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java
index 68befa37..5c5deda8 100644
--- a/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java
+++ b/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java
@@ -29,7 +29,7 @@
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.index.query.QueryBuilders;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.TestGeoShapeFieldMapperPlugin;
@@ -125,7 +125,8 @@ protected boolean supportsQueryName() {
     }
 
     @Override
-    protected void doAssertLuceneQuery(ValidatingLtrQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
+    protected void doAssertLuceneQuery(ValidatingLtrQueryBuilder queryBuilder, Query query,
+                                       SearchExecutionContext context) throws IOException {
         if (StoredFeature.TYPE.equals(queryBuilder.getElement().type())) {
             assertThat(query, instanceOf(MatchNoDocsQuery.class));
         } else if (StoredFeatureSet.TYPE.equals(queryBuilder.getElement().type())) {
diff --git a/src/test/java/com/o19s/es/termstat/TermStatQueryBuilderTests.java b/src/test/java/com/o19s/es/termstat/TermStatQueryBuilderTests.java
index a4f98d45..cdb725f1 100644
--- a/src/test/java/com/o19s/es/termstat/TermStatQueryBuilderTests.java
+++ b/src/test/java/com/o19s/es/termstat/TermStatQueryBuilderTests.java
@@ -5,7 +5,7 @@
 import com.o19s.es.ltr.LtrQueryParserPlugin;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.ParsingException;
-import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.TestGeoShapeFieldMapperPlugin;
@@ -69,7 +69,7 @@ public void testMissingExpr() throws Exception {
     }
 
     @Override
-    protected void doAssertLuceneQuery(TermStatQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
+    protected void doAssertLuceneQuery(TermStatQueryBuilder queryBuilder, Query query, SearchExecutionContext context) throws IOException {
         assertThat(query, instanceOf(TermStatQuery.class));
     }
 }