Cleanup

kderusso · Dec 3, 2024 · 4232154 · 4232154
1 parent f219ff3
commit 4232154
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 55 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/query/InferenceQueryBuilderService.java b/server/src/main/java/org/elasticsearch/index/query/InferenceQueryBuilderService.java
@@ -9,6 +9,12 @@
 package org.elasticsearch.index.query;
 
 import org.elasticsearch.common.TriFunction;
+import org.elasticsearch.plugins.PluginsService;
+import org.elasticsearch.plugins.SearchPlugin;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
 
 public class InferenceQueryBuilderService {
     private final TriFunction<String, String, Boolean, AbstractQueryBuilder<?>> defaultInferenceQueryBuilder;
@@ -22,4 +28,30 @@ public AbstractQueryBuilder<?> getDefaultInferenceQueryBuilder(String fieldName,
             ? defaultInferenceQueryBuilder.apply(fieldName, query, throwOnUnsupportedQueries)
             : null;
     }
+
+    public static InferenceQueryBuilderService build(PluginsService pluginsService) {
+        Objects.requireNonNull(pluginsService);
+
+        List<TriFunction<String, String, Boolean, AbstractQueryBuilder<?>>> definedInferenceQueryBuilders = new ArrayList<>();
+
+        List<SearchPlugin> searchPlugins = pluginsService.filterPlugins(SearchPlugin.class).toList();
+        for (SearchPlugin searchPlugin : searchPlugins) {
+            if (searchPlugin.getDefaultInferenceQueryBuilder() != null) {
+                definedInferenceQueryBuilders.add(searchPlugin.getDefaultInferenceQueryBuilder());
+            }
+        }
+
+        if (definedInferenceQueryBuilders.isEmpty()) {
+            // Backwards compatibility
+            return new InferenceQueryBuilderService(null);
+        }
+
+        if (definedInferenceQueryBuilders.size() != 1) {
+            throw new IllegalStateException(
+                "Expected a single default inference query builder, but found " + definedInferenceQueryBuilders.size()
+            );
+        }
+
+        return new InferenceQueryBuilderService(definedInferenceQueryBuilders.getFirst());
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/index/query/InferenceQueryBuilderServiceBuilder.java b/server/src/main/java/org/elasticsearch/index/query/InferenceQueryBuilderServiceBuilder.java
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java
@@ -108,7 +108,6 @@
 import org.elasticsearch.index.query.CoordinatorRewriteContextProvider;
 import org.elasticsearch.index.query.DataRewriteContext;
 import org.elasticsearch.index.query.InferenceQueryBuilderService;
-import org.elasticsearch.index.query.InferenceQueryBuilderServiceBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.recovery.RecoveryStats;
@@ -382,7 +381,7 @@ public void onRemoval(ShardId shardId, String fieldName, boolean wasEvicted, lon
         this.timestampFieldMapperService = new TimestampFieldMapperService(settings, threadPool, this);
         this.postRecoveryMerger = new PostRecoveryMerger(settings, threadPool.executor(ThreadPool.Names.FORCE_MERGE), this::getShardOrNull);
         this.searchOperationListeners = builder.searchOperationListener;
-        this.inferenceQueryBuilderService = new InferenceQueryBuilderServiceBuilder().pluginsService(pluginsService).build();
+        this.inferenceQueryBuilderService = InferenceQueryBuilderService.build(pluginsService);
     }
 
     private static final String DANGLING_INDICES_UPDATE_THREAD_NAME = "DanglingIndices#updateTask";