Resolving merge conflicts.

README.md

@@ -45,7 +45,7 @@ curl -s http://localhost:9200/.awesome_events/_search -H 'Content-Type: applicat
 Get queries:
 
 ```
-curl -s http://localhost:9200/.awesome_queries/_search | jq
+curl -s http://localhost:9200/.awesome_queries/_search -H "X-ubi-store: awesome" | jq
 ```
 
 Delete the store:

src/main/java/org/opensearch/ubl/HeaderConstants.java

@@ -0,0 +1,15 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.ubl;
+
+public class HeaderConstants {
+
+    public static final String EVENT_STORE_HEADER = "X-ubi-store";
+
+}

src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java

@@ -25,6 +25,7 @@
 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;
@@ -50,6 +51,16 @@ public class UserBehaviorLoggingPlugin extends Plugin implements ActionPlugin {
     private Backend backend;
     private ActionFilter userBehaviorLoggingFilter;
 
+    @Override
+    public Collection<RestHeaderDefinition> getRestHeaders() {
+        return List.of(new RestHeaderDefinition(HeaderConstants.EVENT_STORE_HEADER, false));
+    }
+
+    @Override
+    public Collection<String> getTaskHeaders() {
+        return List.of(HeaderConstants.EVENT_STORE_HEADER);
+    }
+
     @Override
     public List<RestHandler> getRestHandlers(final Settings settings,
                                              final RestController restController,

src/main/java/org/opensearch/ubl/action/UserBehaviorLoggingActionFilter.java

@@ -16,16 +16,18 @@
 import org.opensearch.action.support.ActionFilter;
 import org.opensearch.action.support.ActionFilterChain;
 import org.opensearch.common.settings.Settings;
-import org.opensearch.common.util.concurrent.ThreadContext;
 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.model.QueryResponse;
+import org.opensearch.ubl.HeaderConstants;
 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 UserBehaviorLoggingActionFilter implements ActionFilter {
 
@@ -65,13 +67,18 @@ public void onResponse(Response response) {
 
                 final long startTime = System.currentTimeMillis();
 
-                // Get the search itself.
-                final SearchRequest searchRequest = (SearchRequest) request;
+                final String eventStore = task.getHeader(HeaderConstants.EVENT_STORE_HEADER);
 
-                // Restrict this to only searches of certain indices specified in the settings.
-                //final List<String> indices = Arrays.asList(searchRequest.indices());
-                //final Set<String> indicesToLog = new HashSet<>(Arrays.asList(settings.get(SettingsConstants.INDEX_NAMES).split(",")));
-                //if(indicesToLog.containsAll(indices)) {
+                // If there is no event store header we should not continue anything.
+                if(eventStore != null && !eventStore.trim().isEmpty()) {
+
+                    // Get the search itself.
+                    final SearchRequest searchRequest = (SearchRequest) request;
+
+                    // TODO: Restrict logging to only queries of certain indices specified in the settings.
+                    //final List<String> indices = Arrays.asList(searchRequest.indices());
+                    //final Set<String> indicesToLog = new HashSet<>(Arrays.asList(settings.get(SettingsConstants.INDEX_NAMES).split(",")));
+                    //if(indicesToLog.containsAll(indices)) {
 
                     // Get all search hits from the response.
                     if (response instanceof SearchResponse) {
@@ -97,8 +104,7 @@ public void onResponse(Response response) {
                         try {
 
                             // Persist the query to the backend.
-                            // TODO: How do we know which storeName?
-                            backend.persistQuery("awesome",
+                            backend.persistQuery(eventStore,
                                     new QueryRequest(queryId, query),
                                     new QueryResponse(queryId, queryResponseId, queryResponseHitIds));
 
@@ -111,10 +117,12 @@ public void onResponse(Response response) {
 
                     }
 
-                //}
+                    //}
+
+                    final long elapsedTime = System.currentTimeMillis() - startTime;
+                    LOGGER.info("UBL search request filter took {} ms", elapsedTime);
 
-                final long elapsedTime = System.currentTimeMillis() - startTime;
-                LOGGER.info("UBL search request filter took {} ms", elapsedTime);
+                }
 
                 listener.onResponse(response);