Merge pull request #43 from o19s/query_id

Adding query_id to the response headers
o19s · Feb 19, 2024 · f26f334 · f26f334
2 parents 3352880 + 120170e
commit f26f334
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 23 deletions.
diff --git a/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java b/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java
@@ -10,12 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.rest.RestHeaderDefinition;
-import org.opensearch.ubl.action.UserBehaviorLoggingRestHandler;
-import org.opensearch.ubl.action.UserBehaviorLoggingSearchFilter;
 import org.opensearch.action.support.ActionFilter;
-import org.opensearch.ubl.backends.Backend;
-import org.opensearch.ubl.backends.OpenSearchBackend;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodes;
@@ -25,14 +20,19 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
-import org.opensearch.ubl.events.EventManager;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.rest.RestHeaderDefinition;
 import org.opensearch.script.ScriptService;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.ubl.action.UserBehaviorLoggingActionFilter;
+import org.opensearch.ubl.action.UserBehaviorLoggingRestHandler;
+import org.opensearch.ubl.backends.Backend;
+import org.opensearch.ubl.backends.OpenSearchBackend;
+import org.opensearch.ubl.events.EventManager;
 import org.opensearch.watcher.ResourceWatcherService;
 
 import java.util.ArrayList;
@@ -109,7 +109,7 @@ public Collection<Object> createComponents(
     ) {
 
         this.backend = new OpenSearchBackend(client);
-        this.userBehaviorLoggingFilter =  new UserBehaviorLoggingSearchFilter(backend, environment.settings());
+        this.userBehaviorLoggingFilter =  new UserBehaviorLoggingActionFilter(backend, environment.settings(), threadPool);
 
         LOGGER.info("Creating scheduled task");
 

diff --git a/...tion/UserBehaviorLoggingSearchFilter.java → ...tion/UserBehaviorLoggingActionFilter.java b/...tion/UserBehaviorLoggingSearchFilter.java → ...tion/UserBehaviorLoggingActionFilter.java
@@ -10,7 +10,6 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.lucene.queryparser.flexible.core.util.StringUtils;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -19,24 +18,29 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
+import org.opensearch.tasks.Task;
+import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.ubl.HeaderConstants;
-import org.opensearch.ubl.model.QueryResponse;
 import org.opensearch.ubl.backends.Backend;
 import org.opensearch.ubl.model.QueryRequest;
-import org.opensearch.tasks.Task;
+import org.opensearch.ubl.model.QueryResponse;
 
-import java.util.*;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.UUID;
 
-public class UserBehaviorLoggingSearchFilter implements ActionFilter {
+public class UserBehaviorLoggingActionFilter implements ActionFilter {
 
-    private static final Logger LOGGER = LogManager.getLogger(UserBehaviorLoggingSearchFilter.class);
+    private static final Logger LOGGER = LogManager.getLogger(UserBehaviorLoggingActionFilter.class);
 
     private final Backend backend;
     private final Settings settings;
+    private final ThreadPool threadPool;
 
-    public UserBehaviorLoggingSearchFilter(final Backend backend, final Settings settings) {
+    public UserBehaviorLoggingActionFilter(final Backend backend, final Settings settings, ThreadPool threadPool) {
         this.backend = backend;
         this.settings = settings;
+        this.threadPool = threadPool;
     }
 
     @Override
@@ -59,6 +63,8 @@ public <Request extends ActionRequest, Response extends ActionResponse> void app
             @Override
             public void onResponse(Response response) {
 
+                LOGGER.info("Query ID header: " + task.getHeader("query-id"));
+
                 final long startTime = System.currentTimeMillis();
 
                 final String eventStore = task.getHeader(HeaderConstants.EVENT_STORE_HEADER);
@@ -74,19 +80,19 @@ public void onResponse(Response response) {
                     //final Set<String> indicesToLog = new HashSet<>(Arrays.asList(settings.get(SettingsConstants.INDEX_NAMES).split(",")));
                     //if(indicesToLog.containsAll(indices)) {
 
-                    // Create a UUID for this search request.
-                    final String queryId = UUID.randomUUID().toString();
+                    // Get all search hits from the response.
+                    if (response instanceof SearchResponse) {
 
-                    // The query will be empty when there is no query, e.g. /_search
-                    final String query = searchRequest.source().toString();
+                        // Create a UUID for this search request.
+                        final String queryId = UUID.randomUUID().toString();
 
-                    // Create a UUID for this search response.
-                    final String queryResponseId = UUID.randomUUID().toString();
+                        // The query will be empty when there is no query, e.g. /_search
+                        final String query = searchRequest.source().toString();
 
-                    final List<String> queryResponseHitIds = new LinkedList<>();
+                        // Create a UUID for this search response.
+                        final String queryResponseId = UUID.randomUUID().toString();
 
-                    // Get all search hits from the response.
-                    if (response instanceof SearchResponse) {
+                        final List<String> queryResponseHitIds = new LinkedList<>();
 
                         final SearchResponse searchResponse = (SearchResponse) response;
 
@@ -107,6 +113,8 @@ public void onResponse(Response response) {
                             LOGGER.error("Unable to persist query.", ex);
                         }
 
+                        threadPool.getThreadContext().addResponseHeader("query_id", queryId);
+
                     }
 
                     //}