Routes to pagination or non-pagination execution based on fetch_size parameter.
*/
private Protocol buildProtocolForDefaultQuery(Client client, DefaultQueryAction queryAction)
throws SqlParseException {
- PointInTimeHandler pit = null;
- SearchResponse response;
- SqlOpenSearchRequestBuilder sqlOpenSearchRequestBuilder = queryAction.explain();
+ queryAction.explain();
- pit = new PointInTimeHandlerImpl(client, queryAction.getSelect().getIndexArr());
+ Integer fetchSize = queryAction.getSqlRequest().fetchSize();
+ if (fetchSize != null && fetchSize > 0) {
+ return buildProtocolWithPagination(client, queryAction, fetchSize);
+ } else {
+ return buildProtocolWithoutPagination(client, queryAction);
+ }
+ }
+
+ /** Executes query with pagination support using Point-in-Time (PIT). */
+ private Protocol buildProtocolWithPagination(
+ Client client, DefaultQueryAction queryAction, Integer fetchSize) {
+
+ PointInTimeHandler pit =
+ new PointInTimeHandlerImpl(client, queryAction.getSelect().getIndexArr());
pit.create();
+
+ try {
+ SearchRequestBuilder searchRequest = queryAction.getRequestBuilder();
+ searchRequest.setPointInTime(new PointInTimeBuilder(pit.getPitId()));
+ SearchResponse response = searchRequest.get();
+
+ if (shouldCreateCursor(response, queryAction, fetchSize)) {
+ DefaultCursor cursor = createCursorWithPit(pit, response, queryAction, fetchSize);
+ return new Protocol(client, queryAction, response.getHits(), format, cursor);
+ } else {
+ pit.delete();
+ return new Protocol(client, queryAction, response.getHits(), format, Cursor.NULL_CURSOR);
+ }
+ } catch (RuntimeException e) {
+ try {
+ pit.delete();
+ } catch (RuntimeException deleteException) {
+ LOG.error("Failed to delete PIT", deleteException);
+ Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+ }
+ throw e;
+ }
+ }
+
+ private Protocol buildProtocolWithoutPagination(Client client, DefaultQueryAction queryAction) {
SearchRequestBuilder searchRequest = queryAction.getRequestBuilder();
- searchRequest.setPointInTime(new PointInTimeBuilder(pit.getPitId()));
- response = searchRequest.get();
+ SearchResponse response = searchRequest.get();
+ return new Protocol(client, queryAction, response.getHits(), format, Cursor.NULL_CURSOR);
+ }
- Protocol protocol;
- if (isDefaultCursor(response, queryAction)) {
- DefaultCursor defaultCursor = new DefaultCursor();
- defaultCursor.setLimit(queryAction.getSelect().getRowCount());
- defaultCursor.setFetchSize(queryAction.getSqlRequest().fetchSize());
-
- defaultCursor.setPitId(pit.getPitId());
- defaultCursor.setSearchSourceBuilder(queryAction.getRequestBuilder().request().source());
- defaultCursor.setSortFields(
+ private DefaultCursor createCursorWithPit(
+ PointInTimeHandler pit,
+ SearchResponse response,
+ DefaultQueryAction queryAction,
+ Integer fetchSize) {
+ DefaultCursor cursor = new DefaultCursor();
+ cursor.setLimit(queryAction.getSelect().getRowCount());
+ cursor.setFetchSize(fetchSize);
+ cursor.setPitId(pit.getPitId());
+ cursor.setSearchSourceBuilder(queryAction.getRequestBuilder().request().source());
+
+ if (response.getHits().getHits().length > 0) {
+ cursor.setSortFields(
response.getHits().getAt(response.getHits().getHits().length - 1).getSortValues());
-
- protocol = new Protocol(client, queryAction, response.getHits(), format, defaultCursor);
- } else {
- protocol = new Protocol(client, queryAction, response.getHits(), format, Cursor.NULL_CURSOR);
}
- return protocol;
+ return cursor;
}
- protected boolean isDefaultCursor(SearchResponse searchResponse, DefaultQueryAction queryAction) {
- return queryAction.getSqlRequest().fetchSize() != 0
- && Objects.requireNonNull(searchResponse.getHits().getTotalHits()).value()
- >= queryAction.getSqlRequest().fetchSize();
+ protected boolean shouldCreateCursor(
+ SearchResponse searchResponse, DefaultQueryAction queryAction, Integer fetchSize) {
+ return fetchSize != null
+ && searchResponse.getHits() != null
+ && Objects.requireNonNull(searchResponse.getHits().getTotalHits()).value() >= fetchSize;
}
}
diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutorPitTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutorPitTest.java
new file mode 100644
index 00000000000..ac9178ec5d5
--- /dev/null
+++ b/legacy/src/test/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutorPitTest.java
@@ -0,0 +1,127 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy.executor.format;
+
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.lucene.search.TotalHits;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.opensearch.action.search.SearchRequestBuilder;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.search.SearchHit;
+import org.opensearch.search.SearchHits;
+import org.opensearch.search.builder.PointInTimeBuilder;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.query.DefaultQueryAction;
+import org.opensearch.sql.legacy.query.SqlOpenSearchRequestBuilder;
+import org.opensearch.sql.legacy.request.SqlRequest;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
+import org.opensearch.transport.client.Client;
+
+/** Unit tests for PIT lifecycle management in PrettyFormatRestExecutor. */
+@RunWith(MockitoJUnitRunner.Silent.class)
+public class PrettyFormatRestExecutorPitTest {
+
+ @Mock private Client client;
+ @Mock private DefaultQueryAction queryAction;
+ @Mock private SqlRequest sqlRequest;
+ @Mock private SqlOpenSearchRequestBuilder requestBuilder;
+ @Mock private SearchRequestBuilder searchRequestBuilder;
+ @Mock private SearchResponse searchResponse;
+ @Mock private SearchHit searchHit;
+
+ private PrettyFormatRestExecutor executor;
+ private Map For PIT lifecycle tests, see {@link PrettyFormatRestExecutorPitTest}.
+ */
@RunWith(MockitoJUnitRunner.class)
public class PrettyFormatRestExecutorTest {
@Mock private SearchResponse searchResponse;
@Mock private SearchHit searchHit;
@Mock private DefaultQueryAction queryAction;
- @Mock private SqlRequest sqlRequest;
private PrettyFormatRestExecutor executor;
@Before
public void setUp() {
OpenSearchSettings settings = mock(OpenSearchSettings.class);
LocalClusterState.state().setPluginSettings(settings);
- when(queryAction.getSqlRequest()).thenReturn(sqlRequest);
executor = new PrettyFormatRestExecutor("jdbc");
}
@Test
- public void testIsDefaultCursor_fetchSizeZero() {
- when(sqlRequest.fetchSize()).thenReturn(0);
-
- assertFalse(executor.isDefaultCursor(searchResponse, queryAction));
+ public void testShouldCreateCursor_fetchSizeZero() {
+ assertFalse(executor.shouldCreateCursor(searchResponse, queryAction, 0));
}
@Test
- public void testIsDefaultCursor_totalHitsLessThanFetchSize() {
- when(sqlRequest.fetchSize()).thenReturn(10);
+ public void testShouldCreateCursor_totalHitsLessThanFetchSize() {
when(searchResponse.getHits())
.thenReturn(
new SearchHits(
new SearchHit[] {searchHit}, new TotalHits(5, TotalHits.Relation.EQUAL_TO), 1.0F));
- assertFalse(executor.isDefaultCursor(searchResponse, queryAction));
+ assertFalse(executor.shouldCreateCursor(searchResponse, queryAction, 10));
}
@Test
- public void testIsDefaultCursor_totalHitsGreaterThanOrEqualToFetchSize() {
- when(sqlRequest.fetchSize()).thenReturn(5);
+ public void testShouldCreateCursor_totalHitsGreaterThanOrEqualToFetchSize() {
when(searchResponse.getHits())
.thenReturn(
new SearchHits(
new SearchHit[] {searchHit}, new TotalHits(5, TotalHits.Relation.EQUAL_TO), 1.0F));
- assertTrue(executor.isDefaultCursor(searchResponse, queryAction));
+ assertTrue(executor.shouldCreateCursor(searchResponse, queryAction, 5));
}
}
diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java
index 7b1e2dec0f7..3f465b96d2b 100644
--- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java
+++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java
@@ -130,7 +130,7 @@ void testGetStorageEngineWithWrongAuthType() {
.thenReturn(Collections.emptyList());
PrometheusStorageFactory prometheusStorageFactory = new PrometheusStorageFactory(settings);
HashMap