From 48928d9b45c0aa7914ea582c3019e5df8e236baf Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 9 Aug 2023 22:10:10 +0000 Subject: [PATCH] (#1536) Refactor OpenSearchQueryRequest and move includes to builder (#1937) * #1536: Refactor OpenSearchQueryRequest and move includes to builder (#320) * #1536: Refactor OpenSearchQueryRequest and move incldues to builder Signed-off-by: acarbonetto * #1536: Checkstyle fixes Signed-off-by: acarbonetto * #1536: Checkstyle fixes Signed-off-by: acarbonetto --------- Signed-off-by: acarbonetto * #1536: Spotless Apply Signed-off-by: acarbonetto --------- Signed-off-by: acarbonetto (cherry picked from commit 7d23e0f4b28aa59578bf362e4cebd2eaea693f37) Signed-off-by: github-actions[bot] --- .../request/OpenSearchQueryRequest.java | 24 +++--- .../opensearch/request/OpenSearchRequest.java | 1 + .../request/OpenSearchRequestBuilder.java | 11 ++- .../request/OpenSearchScrollRequest.java | 10 +-- .../client/OpenSearchNodeClientTest.java | 8 +- .../client/OpenSearchRestClientTest.java | 12 +-- .../request/OpenSearchQueryRequestTest.java | 22 +++--- .../request/OpenSearchRequestBuilderTest.java | 76 ++++++++++++++++++- .../request/OpenSearchScrollRequestTest.java | 49 +++--------- .../storage/scan/OpenSearchIndexScanTest.java | 8 +- 10 files changed, 137 insertions(+), 84 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 6da34dd678..919596eee2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -7,7 +7,6 @@ package org.opensearch.sql.opensearch.request; import java.io.IOException; -import java.util.Arrays; import java.util.List; import java.util.function.Consumer; import java.util.function.Function; @@ -20,7 +19,6 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.search.fetch.subphase.FetchSourceContext; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -52,6 +50,14 @@ public class OpenSearchQueryRequest implements OpenSearchRequest { @ToString.Exclude private final OpenSearchExprValueFactory exprValueFactory; + + /** + * List of includes expected in the response. + */ + @EqualsAndHashCode.Exclude + @ToString.Exclude + private final List includes; + /** * Indicate the search already done. */ @@ -61,40 +67,38 @@ public class OpenSearchQueryRequest implements OpenSearchRequest { * Constructor of OpenSearchQueryRequest. */ public OpenSearchQueryRequest(String indexName, int size, - OpenSearchExprValueFactory factory) { - this(new IndexName(indexName), size, factory); + OpenSearchExprValueFactory factory, List includes) { + this(new IndexName(indexName), size, factory, includes); } /** * Constructor of OpenSearchQueryRequest. */ public OpenSearchQueryRequest(IndexName indexName, int size, - OpenSearchExprValueFactory factory) { + OpenSearchExprValueFactory factory, List includes) { this.indexName = indexName; this.sourceBuilder = new SearchSourceBuilder(); sourceBuilder.from(0); sourceBuilder.size(size); sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT); this.exprValueFactory = factory; + this.includes = includes; } /** * Constructor of OpenSearchQueryRequest. */ public OpenSearchQueryRequest(IndexName indexName, SearchSourceBuilder sourceBuilder, - OpenSearchExprValueFactory factory) { + OpenSearchExprValueFactory factory, List includes) { this.indexName = indexName; this.sourceBuilder = sourceBuilder; this.exprValueFactory = factory; + this.includes = includes; } @Override public OpenSearchResponse search(Function searchAction, Function scrollAction) { - FetchSourceContext fetchSource = this.sourceBuilder.fetchSource(); - List includes = fetchSource != null && fetchSource.includes() != null - ? Arrays.asList(fetchSource.includes()) - : List.of(); if (searchDone) { return new OpenSearchResponse(SearchHits.empty(), exprValueFactory, includes); } else { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index f070ac11db..5c9d0033c1 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -24,6 +24,7 @@ * OpenSearch search request. */ public interface OpenSearchRequest extends Writeable { + /** * Default query timeout in minutes. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index bec133f834..80259f15d3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -14,6 +14,7 @@ import static org.opensearch.search.sort.SortOrder.ASC; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.List; import java.util.Map; @@ -98,15 +99,19 @@ public OpenSearchRequestBuilder(int requestedTotalSize, public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, int maxResultWindow, TimeValue scrollTimeout) { int size = requestedTotalSize; + FetchSourceContext fetchSource = this.sourceBuilder.fetchSource(); + List includes = fetchSource != null + ? Arrays.asList(fetchSource.includes()) + : List.of(); if (pageSize == null) { if (startFrom + size > maxResultWindow) { sourceBuilder.size(maxResultWindow - startFrom); return new OpenSearchScrollRequest( - indexName, scrollTimeout, sourceBuilder, exprValueFactory); + indexName, scrollTimeout, sourceBuilder, exprValueFactory, includes); } else { sourceBuilder.from(startFrom); sourceBuilder.size(requestedTotalSize); - return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); + return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory, includes); } } else { if (startFrom != 0) { @@ -114,7 +119,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, } sourceBuilder.size(pageSize); return new OpenSearchScrollRequest(indexName, scrollTimeout, - sourceBuilder, exprValueFactory); + sourceBuilder, exprValueFactory, includes); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 9ffcc42ff7..34e8fcd096 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -7,7 +7,6 @@ package org.opensearch.sql.opensearch.request; import java.io.IOException; -import java.util.Arrays; import java.util.List; import java.util.Objects; import java.util.function.Consumer; @@ -71,13 +70,16 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { private boolean needClean = true; @Getter + @EqualsAndHashCode.Exclude + @ToString.Exclude private final List includes; /** Constructor. */ public OpenSearchScrollRequest(IndexName indexName, TimeValue scrollTimeout, SearchSourceBuilder sourceBuilder, - OpenSearchExprValueFactory exprValueFactory) { + OpenSearchExprValueFactory exprValueFactory, + List includes) { this.indexName = indexName; this.scrollTimeout = scrollTimeout; this.exprValueFactory = exprValueFactory; @@ -86,9 +88,7 @@ public OpenSearchScrollRequest(IndexName indexName, .scroll(scrollTimeout) .source(sourceBuilder); - includes = sourceBuilder.fetchSource() == null - ? List.of() - : Arrays.asList(sourceBuilder.fetchSource().includes()); + this.includes = includes; } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index 71a954bca3..ddf5f27f2a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -324,7 +324,7 @@ void search() { // Verify response for first scroll request OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of("id")); OpenSearchResponse response1 = client.search(request); assertFalse(response1.isEmpty()); @@ -359,7 +359,7 @@ void cleanup() { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); request.setScrollId("scroll123"); // Enforce cleaning by setting a private field. FieldUtils.writeField(request, "needClean", true, true); @@ -376,7 +376,7 @@ void cleanup() { void cleanup_without_scrollId() { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); client.cleanup(request); verify(nodeClient, never()).prepareClearScroll(); } @@ -388,7 +388,7 @@ void cleanup_rethrows_exception() { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); request.setScrollId("scroll123"); // Enforce cleaning by setting a private field. FieldUtils.writeField(request, "needClean", true, true); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index 7f968733c1..409596910e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -307,7 +307,7 @@ void search() throws IOException { // Verify response for first scroll request OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of("id")); OpenSearchResponse response1 = client.search(request); assertFalse(response1.isEmpty()); @@ -329,7 +329,7 @@ void search_with_IOException() throws IOException { IllegalStateException.class, () -> client.search(new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory))); + new SearchSourceBuilder(), factory, List.of()))); } @Test @@ -351,7 +351,7 @@ void scroll_with_IOException() throws IOException { // First request run successfully OpenSearchScrollRequest scrollRequest = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); client.search(scrollRequest); assertThrows( IllegalStateException.class, () -> client.search(scrollRequest)); @@ -370,7 +370,7 @@ void schedule() { void cleanup() { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); // Enforce cleaning by setting a private field. FieldUtils.writeField(request, "needClean", true, true); request.setScrollId("scroll123"); @@ -383,7 +383,7 @@ void cleanup() { void cleanup_without_scrollId() throws IOException { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); client.cleanup(request); verify(restClient, never()).clearScroll(any(), any()); } @@ -395,7 +395,7 @@ void cleanup_with_IOException() { OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + new SearchSourceBuilder(), factory, List.of()); // Enforce cleaning by setting a private field. FieldUtils.writeField(request, "needClean", true, true); request.setScrollId("scroll123"); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index cf548d44f9..b6966f2403 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -18,6 +18,7 @@ import static org.mockito.Mockito.when; import static org.opensearch.sql.opensearch.request.OpenSearchRequest.DEFAULT_QUERY_TIMEOUT; +import java.util.List; import java.util.function.Consumer; import java.util.function.Function; import org.apache.lucene.search.TotalHits; @@ -68,27 +69,25 @@ public class OpenSearchQueryRequestTest { private OpenSearchExprValueFactory factory; private final OpenSearchQueryRequest request = - new OpenSearchQueryRequest("test", 200, factory); + new OpenSearchQueryRequest("test", 200, factory, List.of()); private final OpenSearchQueryRequest remoteRequest = - new OpenSearchQueryRequest("ccs:test", 200, factory); + new OpenSearchQueryRequest("ccs:test", 200, factory, List.of()); @Test void search() { OpenSearchQueryRequest request = new OpenSearchQueryRequest( new OpenSearchRequest.IndexName("test"), sourceBuilder, - factory + factory, + List.of() ); - when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext); - when(fetchSourceContext.includes()).thenReturn(null); when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - verify(fetchSourceContext, times(1)).includes(); assertFalse(searchResponse.isEmpty()); searchResponse = request.search(searchAction, scrollAction); assertTrue(searchResponse.isEmpty()); @@ -100,15 +99,14 @@ void search_withoutContext() { OpenSearchQueryRequest request = new OpenSearchQueryRequest( new OpenSearchRequest.IndexName("test"), sourceBuilder, - factory + factory, + List.of() ); - when(sourceBuilder.fetchSource()).thenReturn(null); when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - verify(sourceBuilder, times(1)).fetchSource(); assertFalse(searchResponse.isEmpty()); assertFalse(request.hasAnotherBatch()); } @@ -118,18 +116,16 @@ void search_withIncludes() { OpenSearchQueryRequest request = new OpenSearchQueryRequest( new OpenSearchRequest.IndexName("test"), sourceBuilder, - factory + factory, + List.of() ); String[] includes = {"_id", "_index"}; - when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext); - when(fetchSourceContext.includes()).thenReturn(includes); when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - verify(fetchSourceContext, times(2)).includes(); assertFalse(searchResponse.isEmpty()); searchResponse = request.search(searchAction, scrollAction); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index e8d15bd0bb..483ea1290e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -99,7 +99,7 @@ void build_query_request() { .size(limit) .timeout(DEFAULT_QUERY_TIMEOUT) .trackScores(true), - exprValueFactory), + exprValueFactory, List.of()), requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } @@ -116,7 +116,7 @@ void build_scroll_request_with_correct_size() { .from(offset) .size(MAX_RESULT_WINDOW - offset) .timeout(DEFAULT_QUERY_TIMEOUT), - exprValueFactory), + exprValueFactory, List.of()), requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } @@ -257,6 +257,78 @@ void test_push_down_project() { .timeout(DEFAULT_QUERY_TIMEOUT) .fetchSource(new String[]{"intA"}, new String[0]), requestBuilder); + + assertEquals( + new OpenSearchQueryRequest( + new OpenSearchRequest.IndexName("test"), + new SearchSourceBuilder() + .from(DEFAULT_OFFSET) + .size(DEFAULT_LIMIT) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource("intA", null), + exprValueFactory, + List.of("intA")), + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); + } + + @Test + void test_push_down_project_limit() { + Set references = Set.of(DSL.ref("intA", INTEGER)); + requestBuilder.pushDownProjects(references); + + Integer limit = 200; + Integer offset = 0; + requestBuilder.pushDownLimit(limit, offset); + + assertSearchSourceBuilder( + new SearchSourceBuilder() + .from(offset) + .size(limit) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource(new String[]{"intA"}, new String[0]), + requestBuilder); + + assertEquals( + new OpenSearchQueryRequest( + new OpenSearchRequest.IndexName("test"), + new SearchSourceBuilder() + .from(offset) + .size(limit) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource("intA", null), + exprValueFactory, + List.of("intA")), + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); + } + + @Test + void test_push_down_project_limit_and_offset() { + Set references = Set.of(DSL.ref("intA", INTEGER)); + requestBuilder.pushDownProjects(references); + + Integer limit = 200; + Integer offset = 10; + requestBuilder.pushDownLimit(limit, offset); + + assertSearchSourceBuilder( + new SearchSourceBuilder() + .from(offset) + .size(limit) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource(new String[]{"intA"}, new String[0]), + requestBuilder); + + assertEquals( + new OpenSearchQueryRequest( + new OpenSearchRequest.IndexName("test"), + new SearchSourceBuilder() + .from(offset) + .size(limit) + .timeout(DEFAULT_QUERY_TIMEOUT) + .fetchSource("intA", null), + exprValueFactory, + List.of("intA")), + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 63c6a5ca7d..4b9233dbc1 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -9,7 +9,6 @@ import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -17,8 +16,6 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.NO_SCROLL_ID; @@ -73,22 +70,13 @@ class OpenSearchScrollRequestTest { private final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); private final OpenSearchScrollRequest request = new OpenSearchScrollRequest( INDEX_NAME, SCROLL_TIMEOUT, - searchSourceBuilder, factory); + searchSourceBuilder, factory, List.of()); @Test void constructor() { - searchSourceBuilder.fetchSource(new String[] {"test"}, null); var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, - searchSourceBuilder, factory); - assertNotEquals(List.of(), request.getIncludes()); - } - - @Test - void constructor2() { - searchSourceBuilder.fetchSource(new String[]{"test"}, null); - var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, - factory); - assertNotEquals(List.of(), request.getIncludes()); + searchSourceBuilder, factory, List.of("test")); + assertEquals(List.of("test"), request.getIncludes()); } @Test @@ -134,7 +122,8 @@ void search() { new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), sourceBuilder, - factory + factory, + List.of() ); when(searchResponse.getHits()).thenReturn(searchHits); @@ -150,14 +139,14 @@ void search_without_context() { new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), sourceBuilder, - factory + factory, + List.of() ); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); OpenSearchResponse response = request.search((sr) -> searchResponse, (sr) -> fail()); - verify(sourceBuilder, times(1)).fetchSource(); assertFalse(response.isEmpty()); } @@ -169,7 +158,8 @@ void search_without_scroll_and_initial_request_should_throw() { new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), sourceBuilder, - factory + factory, + List.of() ); var outStream = new BytesStreamOutput(); request.writeTo(outStream); @@ -193,7 +183,8 @@ void search_withoutIncludes() { new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), sourceBuilder, - factory + factory, + List.of() ); when(searchResponse.getHits()).thenReturn(searchHits); @@ -320,22 +311,4 @@ void setScrollId() { request.setScrollId("test"); assertEquals("test", request.getScrollId()); } - - @Test - void includes() { - - assertIncludes(List.of(), searchSourceBuilder); - - searchSourceBuilder.fetchSource((String[])null, (String[])null); - assertIncludes(List.of(), searchSourceBuilder); - - searchSourceBuilder.fetchSource(new String[] {"test"}, null); - assertIncludes(List.of("test"), searchSourceBuilder); - - } - - void assertIncludes(List expected, SearchSourceBuilder sourceBuilder) { - assertEquals(expected, new OpenSearchScrollRequest( - INDEX_NAME, SCROLL_TIMEOUT, sourceBuilder, factory).getIncludes()); - } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 9e96e3cb6a..9a0957ce75 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -24,6 +24,7 @@ import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; +import java.util.List; import java.util.Map; import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; @@ -113,7 +114,7 @@ void serialize() { when(engine.getClient()).thenReturn(client); when(engine.getTable(any(), any())).thenReturn(index); var request = new OpenSearchScrollRequest( - INDEX_NAME, CURSOR_KEEP_ALIVE, searchSourceBuilder, factory); + INDEX_NAME, CURSOR_KEEP_ALIVE, searchSourceBuilder, factory, List.of()); request.setScrollId("valid-id"); // make a response, so OpenSearchResponse::isEmpty would return true and unset needClean var response = mock(SearchResponse.class); @@ -381,7 +382,7 @@ PushDownAssertion shouldQueryHighlight(QueryBuilder query, HighlightBuilder high .highlighter(highlight) .sort(DOC_FIELD_NAME, ASC); OpenSearchRequest request = - new OpenSearchQueryRequest(EMPLOYEES_INDEX, sourceBuilder, factory); + new OpenSearchQueryRequest(EMPLOYEES_INDEX, sourceBuilder, factory, List.of()); when(client.search(request)).thenReturn(response); var indexScan = new OpenSearchIndexScan(client, @@ -397,7 +398,8 @@ PushDownAssertion shouldQuery(QueryBuilder expected) { .size(QUERY_SIZE) .timeout(CURSOR_KEEP_ALIVE) .sort(DOC_FIELD_NAME, ASC); - OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, builder, factory); + OpenSearchRequest request = + new OpenSearchQueryRequest(EMPLOYEES_INDEX, builder, factory, List.of()); when(client.search(request)).thenReturn(response); var indexScan = new OpenSearchIndexScan(client, 10000, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE));