Adding header to determine appropriate query store

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

@@ -10,6 +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;
@@ -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/UserBehaviorLoggingSearchFilter.java

@@ -10,6 +10,7 @@
 
 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;
@@ -18,6 +19,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
+import org.opensearch.ubl.HeaderConstants;
 import org.opensearch.ubl.model.QueryResponse;
 import org.opensearch.ubl.backends.Backend;
 import org.opensearch.ubl.model.QueryRequest;
@@ -59,13 +61,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)) {
 
                     // Create a UUID for this search request.
                     final String queryId = UUID.randomUUID().toString();
@@ -91,8 +98,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));
 
@@ -103,10 +109,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);