Adding rest and unit test for ubi without a query_id. Requiring ubi b…

…lock to have a query_id. Signed-off-by: jzonthemtn <jzemerick@opensourceconnections.com>
o19s · May 7, 2024 · 1d5d871 · 1d5d871
1 parent 2fb835f
commit 1d5d871
Show file tree

Hide file tree

Showing 4 changed files with 119 additions and 75 deletions.
diff --git a/modules/ubi/src/main/java/org/opensearch/ubi/UbiActionFilter.java b/modules/ubi/src/main/java/org/opensearch/ubi/UbiActionFilter.java
@@ -98,29 +98,34 @@ public void onResponse(Response response) {
                     if (ubiParameters != null) {
 
                         final String queryId = ubiParameters.getQueryId();
-                        final String userQuery = ubiParameters.getUserQuery();
-                        final String userId = ubiParameters.getClientId();
-                        final String objectId = ubiParameters.getObjectId();
 
-                        final List<String> queryResponseHitIds = new LinkedList<>();
+                        if (queryId != null) {
 
-                        for (final SearchHit hit : ((SearchResponse) response).getHits()) {
+                            final String userQuery = ubiParameters.getUserQuery();
+                            final String userId = ubiParameters.getClientId();
+                            final String objectId = ubiParameters.getObjectId();
+
+                            final List<String> queryResponseHitIds = new LinkedList<>();
+
+                            for (final SearchHit hit : ((SearchResponse) response).getHits()) {
+
+                                if (objectId == null || objectId.isEmpty()) {
+                                    // Use the result's docId since no object_id was given for the search.
+                                    queryResponseHitIds.add(String.valueOf(hit.docId()));
+                                } else {
+                                    final Map<String, Object> source = hit.getSourceAsMap();
+                                    queryResponseHitIds.add((String) source.get(objectId));
+                                }
 
-                            if (objectId == null || objectId.isEmpty()) {
-                                // Use the result's docId since no object_id was given for the search.
-                                queryResponseHitIds.add(String.valueOf(hit.docId()));
-                            } else {
-                                final Map<String, Object> source = hit.getSourceAsMap();
-                                queryResponseHitIds.add((String) source.get(objectId));
                             }
 
-                        }
+                            final String queryResponseId = UUID.randomUUID().toString();
+                            final QueryResponse queryResponse = new QueryResponse(queryId, queryResponseId, queryResponseHitIds);
+                            final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, queryResponse);
 
-                        final String queryResponseId = UUID.randomUUID().toString();
-                        final QueryResponse queryResponse = new QueryResponse(queryId, queryResponseId, queryResponseHitIds);
-                        final QueryRequest queryRequest = new QueryRequest(queryId, userQuery, userId, queryResponse);
+                            indexQuery(queryRequest);
 
-                        indexUbiQuery(queryRequest);
+                        }
 
                     }
 
@@ -139,7 +144,7 @@ public void onFailure(Exception ex) {
 
     }
 
-    private void indexUbiQuery(final QueryRequest queryRequest) {
+    private void indexQuery(final QueryRequest queryRequest) {
 
         final IndicesExistsRequest indicesExistsRequest = new IndicesExistsRequest(UBI_EVENTS_INDEX, UBI_QUERIES_INDEX);
 
@@ -148,61 +153,65 @@ private void indexUbiQuery(final QueryRequest queryRequest) {
             @Override
             public void onResponse(IndicesExistsResponse indicesExistsResponse) {
 
-                final Settings indexSettings = Settings.builder()
-                    .put(IndexMetadata.INDEX_NUMBER_OF_SHARDS_SETTING.getKey(), 1)
-                    .put(IndexMetadata.INDEX_AUTO_EXPAND_REPLICAS_SETTING.getKey(), "0-2")
-                    .put(IndexMetadata.SETTING_PRIORITY, Integer.MAX_VALUE)
-                    .build();
+                if (!indicesExistsResponse.isExists()) {
 
-                // Create the UBI events index.
-                final CreateIndexRequest createEventsIndexRequest = new CreateIndexRequest(UBI_EVENTS_INDEX).mapping(
-                    getResourceFile(EVENTS_MAPPING_FILE)
-                ).settings(indexSettings);
+                    final Settings indexSettings = Settings.builder()
+                        .put(IndexMetadata.INDEX_NUMBER_OF_SHARDS_SETTING.getKey(), 1)
+                        .put(IndexMetadata.INDEX_AUTO_EXPAND_REPLICAS_SETTING.getKey(), "0-2")
+                        .put(IndexMetadata.SETTING_PRIORITY, Integer.MAX_VALUE)
+                        .build();
 
-                client.admin().indices().create(createEventsIndexRequest);
+                    // Create the UBI events index.
+                    final CreateIndexRequest createEventsIndexRequest = new CreateIndexRequest(UBI_EVENTS_INDEX).mapping(
+                        getResourceFile(EVENTS_MAPPING_FILE)
+                    ).settings(indexSettings);
 
-                // Create the UBI queries index.
-                final CreateIndexRequest createQueriesIndexRequest = new CreateIndexRequest(UBI_QUERIES_INDEX).mapping(
-                    getResourceFile(QUERIES_MAPPING_FILE)
-                ).settings(indexSettings);
+                    client.admin().indices().create(createEventsIndexRequest);
 
-                client.admin().indices().create(createQueriesIndexRequest);
+                    // Create the UBI queries index.
+                    final CreateIndexRequest createQueriesIndexRequest = new CreateIndexRequest(UBI_QUERIES_INDEX).mapping(
+                        getResourceFile(QUERIES_MAPPING_FILE)
+                    ).settings(indexSettings);
 
-            }
+                    client.admin().indices().create(createQueriesIndexRequest);
 
-            @Override
-            public void onFailure(Exception ex) {
-                LOGGER.error("Error creating UBI indexes.", ex);
-            }
+                }
 
-        });
+                LOGGER.debug(
+                    "Indexing query ID {} with response ID {}",
+                    queryRequest.getQueryId(),
+                    queryRequest.getQueryResponse().getQueryResponseId()
+                );
 
-        LOGGER.debug(
-            "Indexing query ID {} with response ID {}",
-            queryRequest.getQueryId(),
-            queryRequest.getQueryResponse().getQueryResponseId()
-        );
+                // What will be indexed - adheres to the queries-mapping.json
+                final Map<String, Object> source = new HashMap<>();
+                source.put("timestamp", queryRequest.getTimestamp());
+                source.put("query_id", queryRequest.getQueryId());
+                source.put("query_response_id", queryRequest.getQueryResponse().getQueryResponseId());
+                source.put("query_response_object_ids", queryRequest.getQueryResponse().getQueryResponseObjectIds());
+                source.put("user_id", queryRequest.getUserId());
+                source.put("user_query", queryRequest.getUserQuery());
 
-        // What will be indexed - adheres to the queries-mapping.json
-        final Map<String, Object> source = new HashMap<>();
-        source.put("timestamp", queryRequest.getTimestamp());
-        source.put("query_id", queryRequest.getQueryId());
-        source.put("query_response_id", queryRequest.getQueryResponse().getQueryResponseId());
-        source.put("query_response_object_ids", queryRequest.getQueryResponse().getQueryResponseObjectIds());
-        source.put("user_id", queryRequest.getUserId());
-        source.put("user_query", queryRequest.getUserQuery());
+                // Build the index request.
+                final IndexRequest indexRequest = new IndexRequest(UBI_QUERIES_INDEX).source(source, XContentType.JSON);
 
-        // Build the index request.
-        final IndexRequest indexRequest = new IndexRequest(UBI_QUERIES_INDEX).source(source, XContentType.JSON);
+                client.index(indexRequest, new ActionListener<>() {
 
-        client.index(indexRequest, new ActionListener<>() {
+                    @Override
+                    public void onResponse(IndexResponse indexResponse) {}
 
-            @Override
-            public void onResponse(IndexResponse indexResponse) {}
+                    @Override
+                    public void onFailure(Exception e) {
+                        LOGGER.error("Unable to index query into UBI index.", e);
+                    }
+
+                });
+
+            }
 
             @Override
-            public void onFailure(Exception e) {
-                LOGGER.error("Unable to index query into UBI index.", e);
+            public void onFailure(Exception ex) {
+                LOGGER.error("Error creating UBI indexes.", ex);
             }
 
         });

diff --git a/modules/ubi/src/main/java/org/opensearch/ubi/ext/UbiParameters.java b/modules/ubi/src/main/java/org/opensearch/ubi/ext/UbiParameters.java
@@ -22,7 +22,6 @@
 import java.io.IOException;
 import java.util.Objects;
 import java.util.Optional;
-import java.util.UUID;
 
 /**
  * The UBI parameters available in the ext.
@@ -117,7 +116,7 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeString(queryId);
+        out.writeString(getQueryId());
         out.writeOptionalString(userQuery);
         out.writeOptionalString(clientId);
         out.writeOptionalString(objectId);
@@ -159,11 +158,7 @@ public int hashCode() {
      * @return The query ID, or a random UUID if the query ID is <code>null</code>.
      */
     public String getQueryId() {
-        if (queryId == null || queryId.isEmpty()) {
-            return UUID.randomUUID().toString();
-        } else {
-            return queryId;
-        }
+        return queryId;
     }
 
     /**

diff --git a/modules/ubi/src/test/java/org.opensearch.ubi/UbiActionFilterTests.java b/modules/ubi/src/test/java/org.opensearch.ubi/UbiActionFilterTests.java
@@ -9,6 +9,7 @@
 package org.opensearch.ubi;
 
 import org.apache.lucene.search.TotalHits;
+import org.opensearch.action.admin.indices.exists.indices.IndicesExistsRequest;
 import org.opensearch.action.admin.indices.exists.indices.IndicesExistsResponse;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -42,7 +43,7 @@
 public class UbiActionFilterTests extends OpenSearchTestCase {
 
     @SuppressWarnings("unchecked")
-    public void testApplyWithUbi() {
+    public void testApplyWithoutUbiBlock() {
 
         final Client client = mock(Client.class);
         final AdminClient adminClient = mock(AdminClient.class);
@@ -52,7 +53,7 @@ public void testApplyWithUbi() {
         when(adminClient.indices()).thenReturn(indicesAdminClient);
 
         final ActionFuture<IndicesExistsResponse> actionFuture = mock(ActionFuture.class);
-        when(indicesAdminClient.exists(any())).thenReturn(actionFuture);
+        when(indicesAdminClient.exists(any(IndicesExistsRequest.class))).thenReturn(actionFuture);
 
         final UbiActionFilter ubiActionFilter = new UbiActionFilter(client);
         final ActionListener<SearchResponse> listener = mock(ActionListener.class);
@@ -74,14 +75,12 @@ public void testApplyWithUbi() {
             return null;
         }).when(chain).proceed(eq(task), anyString(), eq(request), any());
 
-        final UbiParameters params = new UbiParameters("query_id", "user_query", "client_id", "object_id");
-
         UbiParametersExtBuilder builder = mock(UbiParametersExtBuilder.class);
         final List<SearchExtBuilder> builders = new ArrayList<>();
         builders.add(builder);
 
         when(builder.getWriteableName()).thenReturn(UbiParametersExtBuilder.UBI_PARAMETER_NAME);
-        when(builder.getParams()).thenReturn(params);
+        when(builder.getParams()).thenReturn(null);
 
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         searchSourceBuilder.ext(builders);
@@ -90,12 +89,12 @@ public void testApplyWithUbi() {
 
         ubiActionFilter.apply(task, "ubi", request, listener, chain);
 
-        verify(client).index(any(), any());
+        verify(client, never()).index(any(), any());
 
     }
 
     @SuppressWarnings("unchecked")
-    public void testApplyWithoutUbi() {
+    public void testApplyWithUbiBlockWithoutQueryId() {
 
         final Client client = mock(Client.class);
         final AdminClient adminClient = mock(AdminClient.class);
@@ -105,7 +104,7 @@ public void testApplyWithoutUbi() {
         when(adminClient.indices()).thenReturn(indicesAdminClient);
 
         final ActionFuture<IndicesExistsResponse> actionFuture = mock(ActionFuture.class);
-        when(indicesAdminClient.exists(any())).thenReturn(actionFuture);
+        when(indicesAdminClient.exists(any(IndicesExistsRequest.class))).thenReturn(actionFuture);
 
         final UbiActionFilter ubiActionFilter = new UbiActionFilter(client);
         final ActionListener<SearchResponse> listener = mock(ActionListener.class);
@@ -131,8 +130,10 @@ public void testApplyWithoutUbi() {
         final List<SearchExtBuilder> builders = new ArrayList<>();
         builders.add(builder);
 
+        final UbiParameters ubiParameters = new UbiParameters();
+
         when(builder.getWriteableName()).thenReturn(UbiParametersExtBuilder.UBI_PARAMETER_NAME);
-        when(builder.getParams()).thenReturn(null);
+        when(builder.getParams()).thenReturn(ubiParameters);
 
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         searchSourceBuilder.ext(builders);

diff --git a/...ubi/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/10_queries_without_ubi.yml b/...ubi/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/10_queries_without_ubi.yml
@@ -1,5 +1,5 @@
 ---
-"Query":
+"Query without ubi block":
 
   - do:
       indices.create:
@@ -36,3 +36,42 @@
         index: ubi_queries
 
   - is_false: ''
+
+---
+"Query without query_id":
+
+  - do:
+      indices.create:
+        index: ecommerce
+        body:
+          mappings:
+            { "properties": { "category": { "type": "text" } } }
+
+  - match: { acknowledged: true }
+  - match: { index: "ecommerce"}
+
+  - do:
+      index:
+        index: ecommerce
+        id: 1
+        body: { category: notebook }
+
+  - match: { result: created }
+
+  - do:
+      indices.refresh:
+        index: [ "ecommerce" ]
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: ecommerce
+        body: "{\"query\": {\"match\": {\"category\": \"notebook\"}}, \"ext\": {\"ubi\": {\"user_query\": \"notebook\"}}}"
+
+  - gte: { hits.total: 1 }
+
+  - do:
+      indices.exists:
+        index: ubi_queries
+
+  - is_false: ''