Update to use singleton pattern

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

Author: atris

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

@@ -21,17 +21,66 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 /**
- * Registry for query rewriters.
+ * Registry for query rewriters
  *
  * @opensearch.internal
  */
 public final class QueryRewriterRegistry {
 
     private static final Logger logger = LogManager.getLogger(QueryRewriterRegistry.class);
 
-    private QueryRewriterRegistry() {}
+    private static final QueryRewriterRegistry INSTANCE = new QueryRewriterRegistry();
+
+    /**
+     * Default rewriters.
+     * CopyOnWriteArrayList is used for thread-safety during registration.
+     */
+    private final CopyOnWriteArrayList<QueryRewriter> rewriters;
+
+    private QueryRewriterRegistry() {
+        this.rewriters = new CopyOnWriteArrayList<>();
+
+        // Register default rewriters
+        // Note: TermsMergingRewriter is special - it needs threshold at runtime
+        registerRewriter(new BooleanFlatteningRewriter());
+        registerRewriter(new MustToFilterRewriter());
+        registerRewriter(new MustNotToShouldRewriter());
+        registerRewriter(new MatchAllRemovalRewriter());
+    }
+
+    /**
+     * Get the singleton instance of the registry.
+     */
+    public static QueryRewriterRegistry getInstance() {
+        return INSTANCE;
+    }
+
+    /**
+     * Register a custom query rewriter.
+     *
+     * @param rewriter The rewriter to register
+     */
+    public void registerRewriter(QueryRewriter rewriter) {
+        if (rewriter != null) {
+            rewriters.add(rewriter);
+            logger.info("Registered query rewriter: {}", rewriter.name());
+        }
+    }
+
+    /**
+     * Get a list of all rewriters with the given terms threshold.
+     */
+    private List<QueryRewriter> getRewritersWithThreshold(int termsThreshold) {
+        List<QueryRewriter> allRewriters = new ArrayList<>(rewriters);
+        // Add TermsMergingRewriter with the current threshold
+        // This is added dynamically because it needs the threshold parameter
+        allRewriters.add(new TermsMergingRewriter(termsThreshold));
+        allRewriters.sort(Comparator.comparingInt(QueryRewriter::priority));
+        return allRewriters;
+    }
 
     public static QueryBuilder rewrite(QueryBuilder query, QueryShardContext context, boolean enabled) {
         return rewrite(query, context, enabled, 16);
@@ -42,17 +91,11 @@ public static QueryBuilder rewrite(QueryBuilder query, QueryShardContext context
             return query;
         }
 
-        // Create rewriters with the current threshold
-        List<QueryRewriter> currentRewriters = new ArrayList<>();
-        currentRewriters.add(new BooleanFlatteningRewriter());
-        currentRewriters.add(new MustToFilterRewriter());
-        currentRewriters.add(new MustNotToShouldRewriter());
-        currentRewriters.add(new TermsMergingRewriter(termsThreshold));
-        currentRewriters.add(new MatchAllRemovalRewriter());
-        currentRewriters.sort(Comparator.comparingInt(QueryRewriter::priority));
+        QueryRewriterRegistry registry = getInstance();
+        List<QueryRewriter> sortedRewriters = registry.getRewritersWithThreshold(termsThreshold);
 
         QueryBuilder current = query;
-        for (QueryRewriter rewriter : currentRewriters) {
+        for (QueryRewriter rewriter : sortedRewriters) {
             try {
                 QueryBuilder rewritten = rewriter.rewrite(current, context);
                 if (rewritten != current) {

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

@@ -216,4 +216,85 @@ public void testVeryComplexMixedQuery() {
         assertTrue(result.must().size() >= 1);
         assertTrue(result.filter().size() >= 1);
     }
+
+    public void testCustomRewriterRegistration() {
+        // Create a custom rewriter for testing
+        QueryRewriter customRewriter = new QueryRewriter() {
+            @Override
+            public QueryBuilder rewrite(QueryBuilder query, QueryShardContext context) {
+                if (query instanceof TermQueryBuilder) {
+                    TermQueryBuilder termQuery = (TermQueryBuilder) query;
+                    if ("test_field".equals(termQuery.fieldName()) && "test_value".equals(termQuery.value())) {
+                        // Replace with a different query
+                        return QueryBuilders.termQuery("custom_field", "custom_value");
+                    }
+                } else if (query instanceof BoolQueryBuilder) {
+                    // Recursively apply to nested queries
+                    BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
+                    BoolQueryBuilder rewritten = new BoolQueryBuilder();
+
+                    // Copy settings
+                    rewritten.boost(boolQuery.boost());
+                    rewritten.queryName(boolQuery.queryName());
+                    rewritten.minimumShouldMatch(boolQuery.minimumShouldMatch());
+                    rewritten.adjustPureNegative(boolQuery.adjustPureNegative());
+
+                    // Recursively rewrite clauses
+                    boolean changed = false;
+                    for (QueryBuilder must : boolQuery.must()) {
+                        QueryBuilder rewrittenClause = rewrite(must, context);
+                        rewritten.must(rewrittenClause);
+                        if (rewrittenClause != must) changed = true;
+                    }
+                    for (QueryBuilder filter : boolQuery.filter()) {
+                        QueryBuilder rewrittenClause = rewrite(filter, context);
+                        rewritten.filter(rewrittenClause);
+                        if (rewrittenClause != filter) changed = true;
+                    }
+                    for (QueryBuilder should : boolQuery.should()) {
+                        QueryBuilder rewrittenClause = rewrite(should, context);
+                        rewritten.should(rewrittenClause);
+                        if (rewrittenClause != should) changed = true;
+                    }
+                    for (QueryBuilder mustNot : boolQuery.mustNot()) {
+                        QueryBuilder rewrittenClause = rewrite(mustNot, context);
+                        rewritten.mustNot(rewrittenClause);
+                        if (rewrittenClause != mustNot) changed = true;
+                    }
+
+                    return changed ? rewritten : query;
+                }
+                return query;
+            }
+
+            @Override
+            public int priority() {
+                return 1000; // High priority to ensure it runs last
+            }
+
+            @Override
+            public String name() {
+                return "test_custom_rewriter";
+            }
+        };
+
+        // Register the custom rewriter
+        QueryRewriterRegistry.getInstance().registerRewriter(customRewriter);
+
+        // Test that it's applied
+        QueryBuilder query = QueryBuilders.boolQuery()
+            .must(QueryBuilders.termQuery("test_field", "test_value"))
+            .filter(QueryBuilders.termQuery("other_field", "other_value"));
+
+        QueryBuilder rewritten = QueryRewriterRegistry.rewrite(query, context, true);
+        assertThat(rewritten, instanceOf(BoolQueryBuilder.class));
+        BoolQueryBuilder rewrittenBool = (BoolQueryBuilder) rewritten;
+
+        // The custom rewriter should have replaced the term query
+        assertThat(rewrittenBool.must().size(), equalTo(1));
+        assertThat(rewrittenBool.must().get(0), instanceOf(TermQueryBuilder.class));
+        TermQueryBuilder mustTerm = (TermQueryBuilder) rewrittenBool.must().get(0);
+        assertThat(mustTerm.fieldName(), equalTo("custom_field"));
+        assertThat(mustTerm.value(), equalTo("custom_value"));
+    }
 }