From be5b3a04aedbe1dc36046a2e0d4cb289d87975b3 Mon Sep 17 00:00:00 2001 From: jzonthemtn Date: Tue, 7 May 2024 08:26:34 -0400 Subject: [PATCH] Adding rest and unit test for ubi without a query_id. Requiring ubi block to have a query_id. Signed-off-by: jzonthemtn --- .../org/opensearch/ubi/UbiActionFilter.java | 125 ++++++++++-------- .../org/opensearch/ubi/ext/UbiParameters.java | 9 +- .../UbiActionFilterTests.java | 19 +-- .../_plugins.ubi/10_queries_without_ubi.yml | 41 +++++- 4 files changed, 119 insertions(+), 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 index ff3762dac71fb..cf0f9731a531d 100644 --- 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 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 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 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 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.info( + "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 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 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 index 8f7aa40fba55b..96260b83474d0 100644 --- 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 null. */ 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 index 7f2adc7ed3edf..c3288c1cb6902 100644 --- 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 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 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 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 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 listener = mock(ActionListener.class); @@ -131,8 +130,10 @@ public void testApplyWithoutUbi() { final List 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/modules/ubi/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/10_queries_without_ubi.yml b/modules/ubi/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/10_queries_without_ubi.yml index 8ccc402f58f61..c8ce4149a54c9 100644 --- a/modules/ubi/src/yamlRestTest/resources/rest-api-spec/test/_plugins.ubi/10_queries_without_ubi.yml +++ b/modules/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: ''