Addressed PR comments

Signed-off-by: Guian Gumpac <guian.gumpac@improving.com>
Bit-Quill · Mar 14, 2023 · 3de83d9 · 3de83d9
1 parent 2b3a6db
commit 3de83d9
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 28 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java
@@ -50,24 +50,12 @@ void execute(PhysicalPlan plan, ExecutionContext context,
    * Data class that encapsulates ExprValue.
    */
   @Data
+  @RequiredArgsConstructor
   class QueryResponse {
     private final Schema schema;
     private final List<ExprValue> results;
     private final String rawResponse;
 
-    /**
-     * Constructor for Query Response.
-     *
-     * @param schema  schema of the query
-     * @param results list of expressions
-     * @param rawResponse raw response from OpenSearch server
-     */
-    public QueryResponse(Schema schema, List<ExprValue> results, String rawResponse) {
-      this.schema = schema;
-      this.results = results;
-      this.rawResponse = rawResponse;
-    }
-
     /**
      * Constructor for Query Response.
      *

diff --git a/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java b/core/src/test/java/org/opensearch/sql/executor/QueryServiceTest.java
@@ -134,7 +134,7 @@ Helper executeSuccess(Split split) {
           invocation -> {
             ResponseListener<ExecutionEngine.QueryResponse> listener = invocation.getArgument(2);
             listener.onResponse(
-                new ExecutionEngine.QueryResponse(schema, Collections.emptyList(), ""));
+                new ExecutionEngine.QueryResponse(schema, Collections.emptyList()));
             return null;
           })
           .when(executionEngine)

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java
@@ -34,23 +34,23 @@ public class OpenSearchIndexScan extends TableScanOperator {
   @EqualsAndHashCode.Include
   @Getter
   @ToString.Include
-  protected final OpenSearchRequestBuilder requestBuilder;
+  private final OpenSearchRequestBuilder requestBuilder;
 
   /** Search request. */
   @EqualsAndHashCode.Include
   @ToString.Include
-  protected OpenSearchRequest request;
+  private OpenSearchRequest request;
 
   /** Total query size. */
   @EqualsAndHashCode.Include
   @ToString.Include
-  protected Integer querySize;
+  private Integer querySize;
 
   /** Number of rows returned. */
-  protected Integer queryCount;
+  private Integer queryCount;
 
   /** Search response for current batch. */
-  protected Iterator<ExprValue> iterator;
+  private Iterator<ExprValue> iterator;
 
   @Getter
   private String rawResponse;
@@ -102,7 +102,7 @@ public ExprValue next() {
     return iterator.next();
   }
 
-  private OpenSearchResponse fetchNextBatch() {
+  protected OpenSearchResponse fetchNextBatch() {
     OpenSearchResponse response = client.search(request);
     if (!response.isEmpty()) {
       iterator = response.iterator();

diff --git a/...h/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/...h/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java
@@ -23,7 +23,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
@@ -35,6 +34,7 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.action.search.SearchResponse;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.data.model.ExprValue;
@@ -44,7 +44,6 @@
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector;
-import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
@@ -65,6 +64,10 @@ class OpenSearchExecutionEngineTest {
 
   @Mock private Split split;
 
+  @Mock private static SearchResponse searchResponse;
+
+  @Mock private static OpenSearchExprValueFactory factory;
+
   @BeforeEach
   void setUp() {
     doAnswer(
@@ -212,7 +215,7 @@ public void onFailure(Exception e) {
   }
 
   @Test
-  void executeResponseSuccessfully() {
+  void executeRawResponseSuccessfully() {
     OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector);
     Settings settings = mock(Settings.class);
     when(settings.getSettingValue(QUERY_SIZE_LIMIT)).thenReturn(100);
@@ -297,11 +300,8 @@ public FakeOpenSearchIndexScan(OpenSearchClient client, Settings settings,
     }
 
     @Override
-    public void open() {
-      querySize = requestBuilder.getQuerySize();
-      request = requestBuilder.build();
-      iterator = Collections.emptyIterator();
-      queryCount = 0;
+    protected OpenSearchResponse fetchNextBatch() {
+      return new OpenSearchResponse(searchResponse, factory);
     }
 
     @Override