address comments

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
opensearch-project · Jul 30, 2024 · 4940dcf · 4940dcf
1 parent 31ca5cc
commit 4940dcf
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 28 deletions.
diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupConstants.java b/server/src/main/java/org/opensearch/wlm/QueryGroupConstants.java
diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java
@@ -16,6 +16,8 @@
 import org.opensearch.tasks.CancellableTask;
 
 import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
 
 import static org.opensearch.search.SearchService.NO_TIMEOUT;
 
@@ -25,6 +27,8 @@
 public class QueryGroupTask extends CancellableTask {
 
     private static final Logger logger = LogManager.getLogger(QueryGroupTask.class);
+    public static final String QUERY_GROUP_ID_HEADER = "queryGroupId";
+    public static final Supplier<String> DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP";
     private String queryGroupId;
 
     public QueryGroupTask(long id, String type, String action, String description, TaskId parentTaskId, Map<String, String> headers) {
@@ -47,7 +51,7 @@ public QueryGroupTask(
      * This method should always be called after calling setQueryGroupId at least once on this object
      * @return task queryGroupId
      */
-    public String getQueryGroupId() {
+    public final String getQueryGroupId() {
         if (queryGroupId == null) {
             logger.warn("QueryGroup _id can't be null, It should be set before accessing it. This is abnormal behaviour ");
         }
@@ -59,12 +63,10 @@ public String getQueryGroupId() {
      * This method was defined since the queryGroupId can only be evaluated after task creation
      * @param threadContext current threadContext
      */
-    public void setQueryGroupId(final ThreadContext threadContext) {
-        this.queryGroupId = QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get();
-
-        if (threadContext != null && threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER) != null) {
-            this.queryGroupId = threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER);
-        }
+    public final void setQueryGroupId(final ThreadContext threadContext) {
+        this.queryGroupId = Optional.ofNullable(threadContext)
+            .map(threadContext1 -> threadContext1.getHeader(QUERY_GROUP_ID_HEADER))
+            .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get());
     }
 
     @Override

diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupTaskTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupTaskTests.java
@@ -14,6 +14,9 @@
 
 import java.util.Collections;
 
+import static org.opensearch.wlm.QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER;
+import static org.opensearch.wlm.QueryGroupTask.QUERY_GROUP_ID_HEADER;
+
 public class QueryGroupTaskTests extends OpenSearchTestCase {
     private ThreadPool threadPool;
     private QueryGroupTask sut;
@@ -31,9 +34,9 @@ public void tearDown() throws Exception {
 
     public void testSuccessfulSetQueryGroupId() {
         sut.setQueryGroupId(threadPool.getThreadContext());
-        assertEquals(QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), sut.getQueryGroupId());
+        assertEquals(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), sut.getQueryGroupId());
 
-        threadPool.getThreadContext().putHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER, "akfanglkaglknag2332");
+        threadPool.getThreadContext().putHeader(QUERY_GROUP_ID_HEADER, "akfanglkaglknag2332");
 
         sut.setQueryGroupId(threadPool.getThreadContext());
         assertEquals("akfanglkaglknag2332", sut.getQueryGroupId());