Both id and query is working and basic unittests

Author: srikanthpadakanti

server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java

@@ -65,7 +65,16 @@
 
 import java.io.IOException;
 import java.nio.CharBuffer;
-import java.util.*;
+import java.util.AbstractList;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -427,7 +436,6 @@ public static TermsQueryBuilder fromXContent(XContentParser parser) throws IOExc
         TermsLookup termsLookup = null;
         QueryBuilder nestedQuery = null;
 
-
         String queryName = null;
         float boost = AbstractQueryBuilder.DEFAULT_BOOST;
 
@@ -506,15 +514,9 @@ public static TermsQueryBuilder fromXContent(XContentParser parser) throws IOExc
         }
 
         if (termsLookup != null) {
-            System.out.println("TermsLookup Details:");
-            System.out.println("Index: " + termsLookup.index());
-            System.out.println("ID: " + termsLookup.id());
-            System.out.println("Path: " + termsLookup.path());
-            System.out.println("Query: " + termsLookup.query());
             termsLookup = new TermsLookup(termsLookup.index(), termsLookup.id(), termsLookup.path(), termsLookup.query());
         }
         return new TermsQueryBuilder(fieldName, values, termsLookup).boost(boost).queryName(queryName).valueType(valueType);
-        //return new TermsQueryBuilder(fieldName, values, termsLookup).boost(boost).queryName(queryName).valueType(valueType);
     }
 
     static List<Object> parseValues(XContentParser parser) throws IOException {
@@ -537,17 +539,16 @@ public String getWriteableName() {
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
         if (termsLookup != null && termsLookup.query() != null) {
-            System.out.println("Executing subquery: " + termsLookup.query());
+            // System.out.println("Executing subquery: " + termsLookup.query());
             QueryBuilder rewrittenQuery = termsLookup.query().rewrite(context);
-            //Query query = rewrittenQuery.toQuery(context);
-            //QueryBuilder query = termsLookup.query().rewrite(context);
+            // Query query = rewrittenQuery.toQuery(context);
+            // QueryBuilder query = termsLookup.query().rewrite(context);
 
-            SearchResponse response = context.getClient().search(
-                new SearchRequest(termsLookup.index())
-                    .source(new SearchSourceBuilder().query(rewrittenQuery).fetchSource(true))
-            ).actionGet();
+            SearchResponse response = context.getClient()
+                .search(new SearchRequest(termsLookup.index()).source(new SearchSourceBuilder().query(rewrittenQuery).fetchSource(true)))
+                .actionGet();
 
-            System.out.println("Subquery Response: " + response);
+            // System.out.println("Subquery Response: " + response);
 
             List<Object> terms = new ArrayList<>();
             for (SearchHit hit : response.getHits().getHits()) {
@@ -637,14 +638,15 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
             }
             // Rewrite the subquery using the shard context
             QueryBuilder rewrittenQuery = termsLookup.query().rewrite(shardContext);
-            System.out.println("Rewritten Query: " + rewrittenQuery);
+            // System.out.println("Rewritten Query: " + rewrittenQuery);
 
             SearchResponse response;
             try {
-                response = shardContext.getClient().search(
-                    new SearchRequest(termsLookup.index())
-                        .source(new SearchSourceBuilder().query(rewrittenQuery).fetchSource(true))
-                ).actionGet();
+                response = shardContext.getClient()
+                    .search(
+                        new SearchRequest(termsLookup.index()).source(new SearchSourceBuilder().query(rewrittenQuery).fetchSource(true))
+                    )
+                    .actionGet();
             } catch (Exception e) {
                 throw new IllegalStateException("Failed to execute subquery: " + e.getMessage(), e);
             }
@@ -656,13 +658,14 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
                     try {
                         List<Object> extracted = XContentMapValues.extractRawValues(termsLookup.path(), source);
                         terms.addAll(extracted);
-                        System.out.println("Extracted Terms: " + extracted);
+                        // System.out.println("Extracted Terms: " + extracted);
                     } catch (Exception ex) {
-                        System.err.println("Error extracting path '" + termsLookup.path() + "' from source: " + source);
-                        ex.printStackTrace();
+                        // System.err.println("Error extracting path '" + termsLookup.path() + "' from source: " + source);
+                        throw new IllegalStateException("Failed to execute subquery: " + ex.getMessage(), ex);
                     }
                 } else {
-                    System.err.println("Source is null for hit: " + hit);
+                    // System.err.println("Source is null for hit: " + hit);
+                    throw new IllegalStateException("Source is null for hit: " + hit);
                 }
             }
             // Return a new TermsQueryBuilder with the fetched terms
@@ -705,6 +708,6 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
         }
 
         return this;
-        //return super.doRewrite(queryRewriteContext);
+        // return super.doRewrite(queryRewriteContext);
     }
 }

server/src/main/java/org/opensearch/indices/TermsLookup.java

@@ -69,6 +69,7 @@ public class TermsLookup implements Writeable, ToXContentFragment {
     public TermsLookup(String index, String id, String path) {
         this(index, id, path, null); // Delegate to the existing constructor
     }
+
     public TermsLookup(String index, String id, String path, QueryBuilder query) {
         if (Strings.isEmpty(index)) {
             throw new IllegalArgumentException("index cannot be null or empty for TermsLookup");
@@ -77,7 +78,9 @@ public TermsLookup(String index, String id, String path, QueryBuilder query) {
             throw new IllegalArgumentException("path cannot be null or empty for TermsLookup");
         }
         if (id == null && query == null) {
-            throw new IllegalArgumentException("[" + TermsQueryBuilder.NAME + "] query lookup element requires specifying either the id or the query.");
+            throw new IllegalArgumentException(
+                "[" + TermsQueryBuilder.NAME + "] query lookup element requires specifying either the id or the query."
+            );
         }
         if (id != null && query != null) {
             throw new IllegalArgumentException("[" + TermsQueryBuilder.NAME + "] query lookup element cannot specify both id and query.");
@@ -210,7 +213,7 @@ public String toString() {
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.field("index", index);
-        //builder.field("id", id);
+        // builder.field("id", id);
         if (id != null) {
             builder.field("id", id);
         }