Integration tests for fetch_size, max_result_window, and query.size_l…

…imit (#248) Signed-off-by: MaxKsyunz <maxk@bitquilltech.com>
Bit-Quill · Mar 27, 2023 · a58733e · a58733e
1 parent ca20d16
commit a58733e
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 4 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java
@@ -257,6 +257,17 @@ protected String executeFetchQuery(String query, int fetchSize, String requestTy
     return responseString;
   }
 
+  protected JSONObject executeQueryTemplate(String queryTemplate, String index, int fetchSize)
+      throws IOException {
+    var query = String.format(queryTemplate, index);
+    return new JSONObject(executeFetchQuery(query, fetchSize, "jdbc"));
+  }
+
+  protected JSONObject executeQueryTemplate(String queryTemplate, String index) throws IOException {
+    var query = String.format(queryTemplate, index);
+    return executeQueryTemplate(queryTemplate, index, 4);
+  }
+
   protected String executeFetchLessQuery(String query, String requestType) throws IOException {
 
     String endpoint = "/_plugins/_sql?format=" + requestType;

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationFallbackIT.java
@@ -127,9 +127,5 @@ public void testOrderBy() throws IOException {
     verifyIsV1Cursor(response);
   }
 
-  private JSONObject executeQueryTemplate(String queryTemplate, String index) throws IOException {
-    var query = String.format(queryTemplate, index);
-    return new JSONObject(executeFetchQuery(query, 4, "jdbc"));
-  }
 
 }
diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.sql;
+
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_PHRASE;
+
+import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.After;
+import org.junit.Test;
+import org.opensearch.client.ResponseException;
+import org.opensearch.sql.legacy.SQLIntegTestCase;
+
+public class PaginationWindowIT extends SQLIntegTestCase {
+  @Override
+  public void init() throws IOException {
+    loadIndex(Index.PHRASE);
+  }
+
+  @After
+  void resetParams() throws IOException {
+    resetMaxResultWindow(TEST_INDEX_PHRASE);
+    resetQuerySizeLimit();
+  }
+
+  @Test
+  public void testFetchSizeLessThanMaxResultWindow() throws IOException {
+    setMaxResultWindow(TEST_INDEX_PHRASE, 6);
+    JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5);
+
+    String cursor = "";
+    int numRows = 0;
+    do {
+      // Process response
+      cursor = response.getString("cursor");
+      numRows += response.getJSONArray("datarows").length();
+      response = executeCursorQuery(cursor);
+    } while (response.has("cursor"));
+
+    var countRows = executeJdbcRequest("SELECT COUNT(*) FROM " + TEST_INDEX_PHRASE)
+        .getJSONArray("datarows")
+        .getJSONArray(0)
+        .get(0);
+    assertEquals(countRows, numRows);
+  }
+
+  @Test
+  public void testQuerySizeLimitDoesNotEffectTotalRowsReturned() throws IOException {
+    int querySizeLimit = 4;
+    setQuerySizeLimit(querySizeLimit);
+    JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5);
+    assertTrue(response.getInt("size") > querySizeLimit);
+
+    String cursor = "";
+    int numRows = 0;
+    do {
+      // Process response
+      cursor = response.getString("cursor");
+      numRows += response.getJSONArray("datarows").length();
+      response = executeCursorQuery(cursor);
+    } while (response.has("cursor"));
+
+    var countRows = executeJdbcRequest("SELECT COUNT(*) FROM " + TEST_INDEX_PHRASE)
+        .getJSONArray("datarows")
+        .getJSONArray(0)
+        .get(0);
+    assertEquals(countRows, numRows);
+    assertTrue(numRows > querySizeLimit);
+  }
+
+  @Test
+  public void testQuerySizeLimitDoesNotEffectPageSize() throws IOException {
+    setQuerySizeLimit(3);
+    setMaxResultWindow(TEST_INDEX_PHRASE, 4);
+    var response
+        = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 4);
+    assertEquals(4, response.getInt("size"));
+
+    var response2
+        = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 2);
+    assertEquals(2, response2.getInt("size"));
+  }
+
+  @Test
+  public void testFetchSizeLargerThanResultWindowFails() throws IOException {
+    final int window = 2;
+    setMaxResultWindow(TEST_INDEX_PHRASE, 2);
+    assertThrows(ResponseException.class,
+        () -> executeQueryTemplate("SELECT * FROM %s",
+            TEST_INDEX_PHRASE, window + 1));
+    resetMaxResultWindow(TEST_INDEX_PHRASE);
+  }
+
+
+}