Skip to content

Commit b480916

Browse files
opensearch-trigger-bot[bot]github-actions[bot]
opensearch-trigger-bot[bot]
and
github-actions[bot]
authored
(opensearch-project#1536) Refactor OpenSearchQueryRequest and move includes to builder (opensearch-project#1937) (opensearch-project#1948)
* opensearch-project#1536: Refactor OpenSearchQueryRequest and move includes to builder (#320) * opensearch-project#1536: Refactor OpenSearchQueryRequest and move incldues to builder * opensearch-project#1536: Checkstyle fixes * opensearch-project#1536: Checkstyle fixes --------- * opensearch-project#1536: Spotless Apply --------- (cherry picked from commit 7d23e0f) Signed-off-by: acarbonetto <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent c5b2433 commit b480916

File tree

10 files changed

+137
-84
lines changed

10 files changed

+137
-84
lines changed

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
package org.opensearch.sql.opensearch.request;
88

99
import java.io.IOException;
10-
import java.util.Arrays;
1110
import java.util.List;
1211
import java.util.function.Consumer;
1312
import java.util.function.Function;
@@ -20,7 +19,6 @@
2019
import org.opensearch.core.common.io.stream.StreamOutput;
2120
import org.opensearch.search.SearchHits;
2221
import org.opensearch.search.builder.SearchSourceBuilder;
23-
import org.opensearch.search.fetch.subphase.FetchSourceContext;
2422
import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
2523
import org.opensearch.sql.opensearch.response.OpenSearchResponse;
2624

@@ -52,6 +50,14 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
5250
@ToString.Exclude
5351
private final OpenSearchExprValueFactory exprValueFactory;
5452

53+
54+
/**
55+
* List of includes expected in the response.
56+
*/
57+
@EqualsAndHashCode.Exclude
58+
@ToString.Exclude
59+
private final List<String> includes;
60+
5561
/**
5662
* Indicate the search already done.
5763
*/
@@ -61,40 +67,38 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
6167
* Constructor of OpenSearchQueryRequest.
6268
*/
6369
public OpenSearchQueryRequest(String indexName, int size,
64-
OpenSearchExprValueFactory factory) {
65-
this(new IndexName(indexName), size, factory);
70+
OpenSearchExprValueFactory factory, List<String> includes) {
71+
this(new IndexName(indexName), size, factory, includes);
6672
}
6773

6874
/**
6975
* Constructor of OpenSearchQueryRequest.
7076
*/
7177
public OpenSearchQueryRequest(IndexName indexName, int size,
72-
OpenSearchExprValueFactory factory) {
78+
OpenSearchExprValueFactory factory, List<String> includes) {
7379
this.indexName = indexName;
7480
this.sourceBuilder = new SearchSourceBuilder();
7581
sourceBuilder.from(0);
7682
sourceBuilder.size(size);
7783
sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT);
7884
this.exprValueFactory = factory;
85+
this.includes = includes;
7986
}
8087

8188
/**
8289
* Constructor of OpenSearchQueryRequest.
8390
*/
8491
public OpenSearchQueryRequest(IndexName indexName, SearchSourceBuilder sourceBuilder,
85-
OpenSearchExprValueFactory factory) {
92+
OpenSearchExprValueFactory factory, List<String> includes) {
8693
this.indexName = indexName;
8794
this.sourceBuilder = sourceBuilder;
8895
this.exprValueFactory = factory;
96+
this.includes = includes;
8997
}
9098

9199
@Override
92100
public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
93101
Function<SearchScrollRequest, SearchResponse> scrollAction) {
94-
FetchSourceContext fetchSource = this.sourceBuilder.fetchSource();
95-
List<String> includes = fetchSource != null && fetchSource.includes() != null
96-
? Arrays.asList(fetchSource.includes())
97-
: List.of();
98102
if (searchDone) {
99103
return new OpenSearchResponse(SearchHits.empty(), exprValueFactory, includes);
100104
} else {

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
* OpenSearch search request.
2525
*/
2626
public interface OpenSearchRequest extends Writeable {
27+
2728
/**
2829
* Default query timeout in minutes.
2930
*/

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import static org.opensearch.search.sort.SortOrder.ASC;
1515

1616
import java.util.ArrayList;
17+
import java.util.Arrays;
1718
import java.util.Collection;
1819
import java.util.List;
1920
import java.util.Map;
@@ -98,23 +99,27 @@ public OpenSearchRequestBuilder(int requestedTotalSize,
9899
public OpenSearchRequest build(OpenSearchRequest.IndexName indexName,
99100
int maxResultWindow, TimeValue scrollTimeout) {
100101
int size = requestedTotalSize;
102+
FetchSourceContext fetchSource = this.sourceBuilder.fetchSource();
103+
List<String> includes = fetchSource != null
104+
? Arrays.asList(fetchSource.includes())
105+
: List.of();
101106
if (pageSize == null) {
102107
if (startFrom + size > maxResultWindow) {
103108
sourceBuilder.size(maxResultWindow - startFrom);
104109
return new OpenSearchScrollRequest(
105-
indexName, scrollTimeout, sourceBuilder, exprValueFactory);
110+
indexName, scrollTimeout, sourceBuilder, exprValueFactory, includes);
106111
} else {
107112
sourceBuilder.from(startFrom);
108113
sourceBuilder.size(requestedTotalSize);
109-
return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory);
114+
return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory, includes);
110115
}
111116
} else {
112117
if (startFrom != 0) {
113118
throw new UnsupportedOperationException("Non-zero offset is not supported with pagination");
114119
}
115120
sourceBuilder.size(pageSize);
116121
return new OpenSearchScrollRequest(indexName, scrollTimeout,
117-
sourceBuilder, exprValueFactory);
122+
sourceBuilder, exprValueFactory, includes);
118123
}
119124
}
120125

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
package org.opensearch.sql.opensearch.request;
88

99
import java.io.IOException;
10-
import java.util.Arrays;
1110
import java.util.List;
1211
import java.util.Objects;
1312
import java.util.function.Consumer;
@@ -71,13 +70,16 @@ public class OpenSearchScrollRequest implements OpenSearchRequest {
7170
private boolean needClean = true;
7271

7372
@Getter
73+
@EqualsAndHashCode.Exclude
74+
@ToString.Exclude
7475
private final List<String> includes;
7576

7677
/** Constructor. */
7778
public OpenSearchScrollRequest(IndexName indexName,
7879
TimeValue scrollTimeout,
7980
SearchSourceBuilder sourceBuilder,
80-
OpenSearchExprValueFactory exprValueFactory) {
81+
OpenSearchExprValueFactory exprValueFactory,
82+
List<String> includes) {
8183
this.indexName = indexName;
8284
this.scrollTimeout = scrollTimeout;
8385
this.exprValueFactory = exprValueFactory;
@@ -86,9 +88,7 @@ public OpenSearchScrollRequest(IndexName indexName,
8688
.scroll(scrollTimeout)
8789
.source(sourceBuilder);
8890

89-
includes = sourceBuilder.fetchSource() == null
90-
? List.of()
91-
: Arrays.asList(sourceBuilder.fetchSource().includes());
91+
this.includes = includes;
9292
}
9393

9494

‎opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ void search() {
324324
// Verify response for first scroll request
325325
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
326326
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
327-
new SearchSourceBuilder(), factory);
327+
new SearchSourceBuilder(), factory, List.of("id"));
328328
OpenSearchResponse response1 = client.search(request);
329329
assertFalse(response1.isEmpty());
330330

@@ -359,7 +359,7 @@ void cleanup() {
359359

360360
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
361361
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
362-
new SearchSourceBuilder(), factory);
362+
new SearchSourceBuilder(), factory, List.of());
363363
request.setScrollId("scroll123");
364364
// Enforce cleaning by setting a private field.
365365
FieldUtils.writeField(request, "needClean", true, true);
@@ -376,7 +376,7 @@ void cleanup() {
376376
void cleanup_without_scrollId() {
377377
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
378378
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
379-
new SearchSourceBuilder(), factory);
379+
new SearchSourceBuilder(), factory, List.of());
380380
client.cleanup(request);
381381
verify(nodeClient, never()).prepareClearScroll();
382382
}
@@ -388,7 +388,7 @@ void cleanup_rethrows_exception() {
388388

389389
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
390390
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
391-
new SearchSourceBuilder(), factory);
391+
new SearchSourceBuilder(), factory, List.of());
392392
request.setScrollId("scroll123");
393393
// Enforce cleaning by setting a private field.
394394
FieldUtils.writeField(request, "needClean", true, true);

‎opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,7 @@ void search() throws IOException {
307307
// Verify response for first scroll request
308308
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
309309
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
310-
new SearchSourceBuilder(), factory);
310+
new SearchSourceBuilder(), factory, List.of("id"));
311311
OpenSearchResponse response1 = client.search(request);
312312
assertFalse(response1.isEmpty());
313313

@@ -329,7 +329,7 @@ void search_with_IOException() throws IOException {
329329
IllegalStateException.class,
330330
() -> client.search(new OpenSearchScrollRequest(
331331
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
332-
new SearchSourceBuilder(), factory)));
332+
new SearchSourceBuilder(), factory, List.of())));
333333
}
334334

335335
@Test
@@ -351,7 +351,7 @@ void scroll_with_IOException() throws IOException {
351351
// First request run successfully
352352
OpenSearchScrollRequest scrollRequest = new OpenSearchScrollRequest(
353353
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
354-
new SearchSourceBuilder(), factory);
354+
new SearchSourceBuilder(), factory, List.of());
355355
client.search(scrollRequest);
356356
assertThrows(
357357
IllegalStateException.class, () -> client.search(scrollRequest));
@@ -370,7 +370,7 @@ void schedule() {
370370
void cleanup() {
371371
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
372372
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
373-
new SearchSourceBuilder(), factory);
373+
new SearchSourceBuilder(), factory, List.of());
374374
// Enforce cleaning by setting a private field.
375375
FieldUtils.writeField(request, "needClean", true, true);
376376
request.setScrollId("scroll123");
@@ -383,7 +383,7 @@ void cleanup() {
383383
void cleanup_without_scrollId() throws IOException {
384384
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
385385
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
386-
new SearchSourceBuilder(), factory);
386+
new SearchSourceBuilder(), factory, List.of());
387387
client.cleanup(request);
388388
verify(restClient, never()).clearScroll(any(), any());
389389
}
@@ -395,7 +395,7 @@ void cleanup_with_IOException() {
395395

396396
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
397397
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
398-
new SearchSourceBuilder(), factory);
398+
new SearchSourceBuilder(), factory, List.of());
399399
// Enforce cleaning by setting a private field.
400400
FieldUtils.writeField(request, "needClean", true, true);
401401
request.setScrollId("scroll123");

‎opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import static org.mockito.Mockito.when;
1919
import static org.opensearch.sql.opensearch.request.OpenSearchRequest.DEFAULT_QUERY_TIMEOUT;
2020

21+
import java.util.List;
2122
import java.util.function.Consumer;
2223
import java.util.function.Function;
2324
import org.apache.lucene.search.TotalHits;
@@ -68,27 +69,25 @@ public class OpenSearchQueryRequestTest {
6869
private OpenSearchExprValueFactory factory;
6970

7071
private final OpenSearchQueryRequest request =
71-
new OpenSearchQueryRequest("test", 200, factory);
72+
new OpenSearchQueryRequest("test", 200, factory, List.of());
7273

7374
private final OpenSearchQueryRequest remoteRequest =
74-
new OpenSearchQueryRequest("ccs:test", 200, factory);
75+
new OpenSearchQueryRequest("ccs:test", 200, factory, List.of());
7576

7677
@Test
7778
void search() {
7879
OpenSearchQueryRequest request = new OpenSearchQueryRequest(
7980
new OpenSearchRequest.IndexName("test"),
8081
sourceBuilder,
81-
factory
82+
factory,
83+
List.of()
8284
);
8385

84-
when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext);
85-
when(fetchSourceContext.includes()).thenReturn(null);
8686
when(searchAction.apply(any())).thenReturn(searchResponse);
8787
when(searchResponse.getHits()).thenReturn(searchHits);
8888
when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
8989

9090
OpenSearchResponse searchResponse = request.search(searchAction, scrollAction);
91-
verify(fetchSourceContext, times(1)).includes();
9291
assertFalse(searchResponse.isEmpty());
9392
searchResponse = request.search(searchAction, scrollAction);
9493
assertTrue(searchResponse.isEmpty());
@@ -100,15 +99,14 @@ void search_withoutContext() {
10099
OpenSearchQueryRequest request = new OpenSearchQueryRequest(
101100
new OpenSearchRequest.IndexName("test"),
102101
sourceBuilder,
103-
factory
102+
factory,
103+
List.of()
104104
);
105105

106-
when(sourceBuilder.fetchSource()).thenReturn(null);
107106
when(searchAction.apply(any())).thenReturn(searchResponse);
108107
when(searchResponse.getHits()).thenReturn(searchHits);
109108
when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
110109
OpenSearchResponse searchResponse = request.search(searchAction, scrollAction);
111-
verify(sourceBuilder, times(1)).fetchSource();
112110
assertFalse(searchResponse.isEmpty());
113111
assertFalse(request.hasAnotherBatch());
114112
}
@@ -118,18 +116,16 @@ void search_withIncludes() {
118116
OpenSearchQueryRequest request = new OpenSearchQueryRequest(
119117
new OpenSearchRequest.IndexName("test"),
120118
sourceBuilder,
121-
factory
119+
factory,
120+
List.of()
122121
);
123122

124123
String[] includes = {"_id", "_index"};
125-
when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext);
126-
when(fetchSourceContext.includes()).thenReturn(includes);
127124
when(searchAction.apply(any())).thenReturn(searchResponse);
128125
when(searchResponse.getHits()).thenReturn(searchHits);
129126
when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
130127

131128
OpenSearchResponse searchResponse = request.search(searchAction, scrollAction);
132-
verify(fetchSourceContext, times(2)).includes();
133129
assertFalse(searchResponse.isEmpty());
134130

135131
searchResponse = request.search(searchAction, scrollAction);

0 commit comments

Comments
 (0)