Require query name with no default

kderusso · Dec 11, 2024 · f4c26ae · f4c26ae
1 parent 6a92936
commit f4c26ae
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 20 deletions.
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java b/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java
@@ -245,15 +245,15 @@ public IndicesService build() {
         var queryRewriteInterceptors = pluginsService.filterPlugins(SearchPlugin.class)
             .map(SearchPlugin::getQueryRewriteInterceptors)
             .flatMap(List::stream)
-            .collect(Collectors.toMap(QueryRewriteInterceptor::getName, interceptor -> {
-                if (interceptor.getName() == null) {
+            .collect(Collectors.toMap(QueryRewriteInterceptor::getQueryName, interceptor -> {
+                if (interceptor.getQueryName() == null) {
                     throw new IllegalArgumentException("QueryRewriteInterceptor [" + interceptor.getClass().getName() + "] requires name");
                 }
                 return interceptor;
             }, (a, b) -> {
                 throw new IllegalStateException(
                     "Conflicting rewrite interceptors ["
-                        + a.getName()
+                        + a.getQueryName()
                         + "] found in ["
                         + a.getClass().getName()
                         + "] and ["

diff --git a/...er/src/main/java/org/elasticsearch/plugins/internal/rewriter/QueryRewriteInterceptor.java b/...er/src/main/java/org/elasticsearch/plugins/internal/rewriter/QueryRewriteInterceptor.java
@@ -30,27 +30,27 @@ public interface QueryRewriteInterceptor {
     QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder);
 
     /**
-     * Implementing classes should override this with the name of the query that is being
-     * intercepted.
+     * Name of the query to be intercepted and rewritten.
      */
-    default String getName() {
-        return null;
-    }
+    String getQueryName();
 
     static QueryRewriteInterceptor multi(Map<String, QueryRewriteInterceptor> interceptors) {
-        if (interceptors.isEmpty()) {
-            return (context, queryBuilder) -> queryBuilder;
-        }
-        return new CompositeQueryRewriteInterceptor(interceptors);
+        return interceptors.isEmpty() ? new NoOpQueryRewriteInterceptor() : new CompositeQueryRewriteInterceptor(interceptors);
     }
 
     class CompositeQueryRewriteInterceptor implements QueryRewriteInterceptor {
+        final String NAME = "composite";
         private final Map<String, QueryRewriteInterceptor> interceptors;
 
         private CompositeQueryRewriteInterceptor(Map<String, QueryRewriteInterceptor> interceptors) {
             this.interceptors = interceptors;
         }
 
+        @Override
+        public String getQueryName() {
+            return NAME;
+        }
+
         @Override
         public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) {
             QueryRewriteInterceptor interceptor = interceptors.get(queryBuilder.getName());
@@ -60,4 +60,16 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
             return queryBuilder;
         }
     }
+
+    class NoOpQueryRewriteInterceptor implements QueryRewriteInterceptor {
+        @Override
+        public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) {
+            return queryBuilder;
+        }
+
+        @Override
+        public String getQueryName() {
+            return null;
+        }
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapperTests.java b/server/src/test/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapperTests.java
@@ -13,15 +13,17 @@
 
 public class InterceptedQueryBuilderWrapperTests extends ESTestCase {
 
-    public void testInterceptedQueryBuildersAreUnwrapped() {
+    public void testQueryNameReturnsWrappedQueryBuilder() {
         MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder();
         InterceptedQueryBuilderWrapper interceptedQueryBuilderWrapper = new InterceptedQueryBuilderWrapper(matchAllQueryBuilder);
-        assertNotEquals(matchAllQueryBuilder, interceptedQueryBuilderWrapper);
-        assertEquals(matchAllQueryBuilder, interceptedQueryBuilderWrapper.queryBuilder);
-        InterceptedQueryBuilderWrapper interceptedQueryBuilderWrapper2 = new InterceptedQueryBuilderWrapper(interceptedQueryBuilderWrapper);
-        assertNotEquals(matchAllQueryBuilder, interceptedQueryBuilderWrapper2);
-        assertEquals(interceptedQueryBuilderWrapper, interceptedQueryBuilderWrapper2);
-        assertEquals(matchAllQueryBuilder, interceptedQueryBuilderWrapper2.queryBuilder);
+        QueryBuilder namedQuery = interceptedQueryBuilderWrapper.queryName(randomAlphaOfLengthBetween(5, 10));
+        assertTrue(namedQuery instanceof InterceptedQueryBuilderWrapper);
     }
 
+    public void testQueryBoostReturnsWrappedQueryBuilder() {
+        MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder();
+        InterceptedQueryBuilderWrapper interceptedQueryBuilderWrapper = new InterceptedQueryBuilderWrapper(matchAllQueryBuilder);
+        QueryBuilder boostedQuery = interceptedQueryBuilderWrapper.boost(randomFloat());
+        assertTrue(boostedQuery instanceof InterceptedQueryBuilderWrapper);
+    }
 }
diff --git a/...rc/main/java/org/elasticsearch/plugins/internal/rewriter/MockQueryRewriteInterceptor.java b/...rc/main/java/org/elasticsearch/plugins/internal/rewriter/MockQueryRewriteInterceptor.java
@@ -18,4 +18,9 @@ public class MockQueryRewriteInterceptor implements QueryRewriteInterceptor {
     public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) {
         return queryBuilder;
     }
+
+    @Override
+    public String getQueryName() {
+        return this.getClass().getSimpleName();
+    }
 }
diff --git a/.../java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java b/.../java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java
@@ -75,7 +75,7 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
     }
 
     @Override
-    public String getName() {
+    public String getQueryName() {
         return MatchQueryBuilder.NAME;
     }