Skip to content

Commit 33d0991

Browse files
acarbonettoacarbonetto
committed
opensearch-project#1536: Refactor OpenSearchQueryRequest and move includes to builder (#320)
* opensearch-project#1536: Refactor OpenSearchQueryRequest and move incldues to builder Signed-off-by: acarbonetto <[email protected]> * opensearch-project#1536: Checkstyle fixes Signed-off-by: acarbonetto <[email protected]> * opensearch-project#1536: Checkstyle fixes Signed-off-by: acarbonetto <[email protected]> --------- Signed-off-by: acarbonetto <[email protected]>
1 parent d00dc4d commit 33d0991

File tree

10 files changed

+137
-78
lines changed

10 files changed

+137
-78
lines changed

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,14 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
5252
@ToString.Exclude
5353
private final OpenSearchExprValueFactory exprValueFactory;
5454

55+
56+
/**
57+
* List of includes expected in the response.
58+
*/
59+
@EqualsAndHashCode.Exclude
60+
@ToString.Exclude
61+
private final List<String> includes;
62+
5563
/**
5664
* Indicate the search already done.
5765
*/
@@ -61,40 +69,38 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
6169
* Constructor of OpenSearchQueryRequest.
6270
*/
6371
public OpenSearchQueryRequest(String indexName, int size,
64-
OpenSearchExprValueFactory factory) {
65-
this(new IndexName(indexName), size, factory);
72+
OpenSearchExprValueFactory factory, List<String> includes) {
73+
this(new IndexName(indexName), size, factory, includes);
6674
}
6775

6876
/**
6977
* Constructor of OpenSearchQueryRequest.
7078
*/
7179
public OpenSearchQueryRequest(IndexName indexName, int size,
72-
OpenSearchExprValueFactory factory) {
80+
OpenSearchExprValueFactory factory, List<String> includes) {
7381
this.indexName = indexName;
7482
this.sourceBuilder = new SearchSourceBuilder();
7583
sourceBuilder.from(0);
7684
sourceBuilder.size(size);
7785
sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT);
7886
this.exprValueFactory = factory;
87+
this.includes = includes;
7988
}
8089

8190
/**
8291
* Constructor of OpenSearchQueryRequest.
8392
*/
8493
public OpenSearchQueryRequest(IndexName indexName, SearchSourceBuilder sourceBuilder,
85-
OpenSearchExprValueFactory factory) {
94+
OpenSearchExprValueFactory factory, List<String> includes) {
8695
this.indexName = indexName;
8796
this.sourceBuilder = sourceBuilder;
8897
this.exprValueFactory = factory;
98+
this.includes = includes;
8999
}
90100

91101
@Override
92102
public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
93103
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();
98104
if (searchDone) {
99105
return new OpenSearchResponse(SearchHits.empty(), exprValueFactory, includes);
100106
} 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 & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,13 +71,16 @@ public class OpenSearchScrollRequest implements OpenSearchRequest {
7171
private boolean needClean = true;
7272

7373
@Getter
74+
@EqualsAndHashCode.Exclude
75+
@ToString.Exclude
7476
private final List<String> includes;
7577

7678
/** Constructor. */
7779
public OpenSearchScrollRequest(IndexName indexName,
7880
TimeValue scrollTimeout,
7981
SearchSourceBuilder sourceBuilder,
80-
OpenSearchExprValueFactory exprValueFactory) {
82+
OpenSearchExprValueFactory exprValueFactory,
83+
List<String> includes) {
8184
this.indexName = indexName;
8285
this.scrollTimeout = scrollTimeout;
8386
this.exprValueFactory = exprValueFactory;
@@ -86,9 +89,7 @@ public OpenSearchScrollRequest(IndexName indexName,
8689
.scroll(scrollTimeout)
8790
.source(sourceBuilder);
8891

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

9495

‎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
@@ -323,7 +323,7 @@ void search() {
323323
// Verify response for first scroll request
324324
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
325325
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
326-
new SearchSourceBuilder(), factory);
326+
new SearchSourceBuilder(), factory, List.of("id"));
327327
OpenSearchResponse response1 = client.search(request);
328328
assertFalse(response1.isEmpty());
329329

@@ -358,7 +358,7 @@ void cleanup() {
358358

359359
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
360360
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
361-
new SearchSourceBuilder(), factory);
361+
new SearchSourceBuilder(), factory, List.of());
362362
request.setScrollId("scroll123");
363363
// Enforce cleaning by setting a private field.
364364
FieldUtils.writeField(request, "needClean", true, true);
@@ -375,7 +375,7 @@ void cleanup() {
375375
void cleanup_without_scrollId() {
376376
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
377377
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
378-
new SearchSourceBuilder(), factory);
378+
new SearchSourceBuilder(), factory, List.of());
379379
client.cleanup(request);
380380
verify(nodeClient, never()).prepareClearScroll();
381381
}
@@ -387,7 +387,7 @@ void cleanup_rethrows_exception() {
387387

388388
OpenSearchScrollRequest request = new OpenSearchScrollRequest(
389389
new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1),
390-
new SearchSourceBuilder(), factory);
390+
new SearchSourceBuilder(), factory, List.of());
391391
request.setScrollId("scroll123");
392392
// Enforce cleaning by setting a private field.
393393
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)