More comments

Signed-off-by: Atri Sharma <atri.jiit@gmail.com>

Author: atris

server/src/main/java/org/opensearch/search/SearchService.java

@@ -1516,10 +1516,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         if (source.query() != null) {
             QueryBuilder query = source.query();
 
-            // Apply query rewriting optimizations if enabled
-            if (QUERY_REWRITING_ENABLED_SETTING.get(clusterService.getSettings())) {
-                query = QueryRewriterRegistry.rewrite(query, queryShardContext, true);
-            }
+            // Apply query rewriting optimizations
+            query = QueryRewriterRegistry.rewrite(query, queryShardContext);
 
             InnerHitContextBuilder.extractInnerHits(query, innerHitBuilders);
             context.parsedQuery(queryShardContext.toQuery(query));

server/src/main/java/org/opensearch/search/query/QueryRewriterRegistry.java

@@ -14,6 +14,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.search.SearchService;
 import org.opensearch.search.query.rewriters.BooleanFlatteningRewriter;
 import org.opensearch.search.query.rewriters.MatchAllRemovalRewriter;
 import org.opensearch.search.query.rewriters.MustNotToShouldRewriter;
@@ -47,6 +48,11 @@ public final class QueryRewriterRegistry {
      */
     private final TermsMergingRewriter termsMergingRewriter;
 
+    /**
+     * Whether query rewriting is enabled.
+     */
+    private volatile boolean enabled;
+
     private QueryRewriterRegistry() {
         this.rewriters = new CopyOnWriteArrayList<>();
 
@@ -89,15 +95,21 @@ public void registerRewriter(QueryRewriter rewriter) {
      * @param clusterSettings Cluster settings for registering update consumers
      */
     public static void initialize(Settings settings, ClusterSettings clusterSettings) {
-        getInstance().termsMergingRewriter.initialize(settings, clusterSettings);
+        QueryRewriterRegistry instance = getInstance();
+        instance.termsMergingRewriter.initialize(settings, clusterSettings);
+        instance.enabled = SearchService.QUERY_REWRITING_ENABLED_SETTING.get(settings);
+        clusterSettings.addSettingsUpdateConsumer(
+            SearchService.QUERY_REWRITING_ENABLED_SETTING,
+            (Boolean enabled) -> instance.enabled = enabled
+        );
     }
 
-    public static QueryBuilder rewrite(QueryBuilder query, QueryShardContext context, boolean enabled) {
-        if (!enabled || query == null) {
+    public static QueryBuilder rewrite(QueryBuilder query, QueryShardContext context) {
+        QueryRewriterRegistry registry = getInstance();
+        if (!registry.enabled || query == null) {
             return query;
         }
 
-        QueryRewriterRegistry registry = getInstance();
         List<QueryRewriter> sortedRewriters = new ArrayList<>(registry.rewriters);
         sortedRewriters.sort(Comparator.comparingInt(QueryRewriter::priority));
 

server/src/test/java/org/opensearch/search/query/QueryRewriterRegistryTests.java

@@ -8,13 +8,16 @@
 
 package org.opensearch.search.query;
 
+import org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.RangeQueryBuilder;
 import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.search.SearchService;
 import org.opensearch.test.OpenSearchTestCase;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -25,6 +28,18 @@ public class QueryRewriterRegistryTests extends OpenSearchTestCase {
 
     private final QueryShardContext context = mock(QueryShardContext.class);
 
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+        // Initialize registry with default settings
+        Settings settings = Settings.builder()
+            .put(SearchService.QUERY_REWRITING_ENABLED_SETTING.getKey(), true)
+            .put(SearchService.QUERY_REWRITING_TERMS_THRESHOLD_SETTING.getKey(), 16)
+            .build();
+        ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        QueryRewriterRegistry.initialize(settings, clusterSettings);
+    }
+
     public void testCompleteRewritingPipeline() {
         // Test that all rewriters work together correctly
         QueryBuilder nestedBool = QueryBuilders.boolQuery()
@@ -38,7 +53,7 @@ public void testCompleteRewritingPipeline() {
             .filter(QueryBuilders.termQuery("type", "product"))
             .filter(QueryBuilders.termQuery("type", "service"));
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
         BoolQueryBuilder rewrittenBool = (BoolQueryBuilder) rewritten;
 
@@ -56,18 +71,31 @@ public void testCompleteRewritingPipeline() {
     }
 
     public void testDisabledRewriting() {
-        // When disabled, query should not be modified
+        // Test disabled rewriting via settings
+        Settings settings = Settings.builder().put(SearchService.QUERY_REWRITING_ENABLED_SETTING.getKey(), false).build();
+        ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+
+        // Initialize with disabled setting
+        QueryRewriterRegistry.initialize(settings, clusterSettings);
+
         QueryBuilder query = QueryBuilders.boolQuery()
             .must(QueryBuilders.matchAllQuery())
             .filter(QueryBuilders.termQuery("field", "value"));
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, false);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertSame(query, rewritten);
+
+        // Enable via settings update
+        clusterSettings.applySettings(Settings.builder().put(SearchService.QUERY_REWRITING_ENABLED_SETTING.getKey(), true).build());
+
+        // Now it should rewrite
+        QueryBuilder rewritten2 = QueryRewriterRegistry.rewrite(query, context);
+        assertNotSame(query, rewritten2);
     }
 
     public void testNullQuery() {
         // Null query should return null
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(null, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(null, context);
         assertNull(rewritten);
     }
 
@@ -82,7 +110,7 @@ public void testRewriterPriorityOrder() {
                     .must(QueryBuilders.termQuery("field", "value2"))
             );
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(deeplyNested, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(deeplyNested, context);
         assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
         BoolQueryBuilder rewrittenBool = (BoolQueryBuilder) rewritten;
 
@@ -113,7 +141,7 @@ public void testComplexRealWorldQuery() {
             .must(QueryBuilders.rangeQuery("price").gte(500).lte(2000))
             .mustNot(QueryBuilders.termQuery("status", "discontinued"));
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
         BoolQueryBuilder rewrittenBool = (BoolQueryBuilder) rewritten;
 
@@ -155,7 +183,7 @@ public void testPerformanceMetrics() {
             .filter(QueryBuilders.matchAllQuery());
 
         // Should not throw any exceptions
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertNotNull(rewritten);
     }
 
@@ -167,7 +195,7 @@ public void testRewriterErrorHandling() {
             .filter(QueryBuilders.matchAllQuery());
 
         // Even if one rewriter fails, others should still be applied
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertNotNull(rewritten);
     }
 
@@ -202,7 +230,7 @@ public void testVeryComplexMixedQuery() {
             .mustNot(QueryBuilders.boolQuery().must(QueryBuilders.termQuery("archived", "true")))
             .minimumShouldMatch(1);
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
         BoolQueryBuilder result = (BoolQueryBuilder) rewritten;
 
@@ -286,7 +314,7 @@ public String name() {
             .must(QueryBuilders.termQuery("test_field", "test_value"))
             .filter(QueryBuilders.termQuery("other_field", "other_value"));
 
-        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context);
         assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
         BoolQueryBuilder rewrittenBool = (BoolQueryBuilder) rewritten;