feat: Add result set constraints and chunk-based polling (#47)

Prevent unnecessary result set iteration:
- Add `DataCloudStatement::setResultSetConstraints()` method to limit
rows and bytes
- Support byte limit configuration via `resultset.bytelimit` property

Chunk-Based result polling:
- Implement waitForChunksAvailable method for pagination

Minor changes:
- Changes default byte size when using row-based access to the server's
largest allowed value
- Implement `getRow()` which is optional on type forward-only
- Add examples demonstrating result set constraints with various styles
of pagination
- Prefer `float4` over `real` since it's more explict

Author: jschneidereit

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/ArrowStreamReaderCursor.java

@@ -25,7 +25,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
-import lombok.AllArgsConstructor;
+import lombok.RequiredArgsConstructor;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
@@ -36,15 +36,18 @@
 import org.apache.calcite.avatica.util.AbstractCursor;
 import org.apache.calcite.avatica.util.ArrayImpl;
 
-@AllArgsConstructor
+@RequiredArgsConstructor
 @Slf4j
 class ArrowStreamReaderCursor extends AbstractCursor {
 
     private static final int INIT_ROW_NUMBER = -1;
 
     private final ArrowStreamReader reader;
 
-    private final AtomicInteger currentRow = new AtomicInteger(INIT_ROW_NUMBER);
+    @lombok.Getter
+    private int rowsSeen = 0;
+
+    private final AtomicInteger currentIndex = new AtomicInteger(INIT_ROW_NUMBER);
 
     private void wasNullConsumer(boolean wasNull) {
         this.wasNull[0] = wasNull;
@@ -65,13 +68,13 @@ public List<Accessor> createAccessors(
     }
 
     private Accessor createAccessor(FieldVector vector) throws SQLException {
-        return QueryJDBCAccessorFactory.createAccessor(vector, currentRow::get, this::wasNullConsumer);
+        return QueryJDBCAccessorFactory.createAccessor(vector, currentIndex::get, this::wasNullConsumer);
     }
 
     private boolean loadNextBatch() throws SQLException {
         try {
             if (reader.loadNextBatch()) {
-                currentRow.set(0);
+                currentIndex.set(0);
                 return true;
             }
         } catch (IOException e) {
@@ -83,11 +86,15 @@ private boolean loadNextBatch() throws SQLException {
     @SneakyThrows
     @Override
     public boolean next() {
-        val current = currentRow.incrementAndGet();
+        val current = currentIndex.incrementAndGet();
         val total = getSchemaRoot().getRowCount();
 
         try {
-            return current < total || loadNextBatch();
+            val next = current < total || loadNextBatch();
+            if (next) {
+                rowsSeen++;
+            }
+            return next;
         } catch (Exception e) {
             throw new DataCloudJDBCException("Failed to load next batch", e);
         }

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudConnection.java

@@ -185,25 +185,44 @@ public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId) t
     }
 
     /**
-     * Checks if all the query's results are ready, the row count and chunk count will be stable.
+     * Checks if a given row range is available for a query.
+     * This method will wait until the row range specified by the other params is available in the given timeout.
+     *
      * @param queryId The identifier of the query to check
      * @param offset The starting row offset.
      * @param limit The quantity of rows relative to the offset to wait for
      * @param timeout The duration to wait for the engine have results produced.
      * @param allowLessThan Whether to return early when the available rows is less than {@code offset + limit}
-     * @return The final {@link DataCloudQueryStatus} the server replied with.
+     * @return The last {@link DataCloudQueryStatus} the server replied with.
      */
     public DataCloudQueryStatus waitForRowsAvailable(
             String queryId, long offset, long limit, Duration timeout, boolean allowLessThan)
             throws DataCloudJDBCException {
         return executor.waitForRowsAvailable(queryId, offset, limit, timeout, allowLessThan);
     }
 
+    /**
+     * Checks if a given chunk range is available for a query.
+     * This method will wait until the chunk range specified by the other params is available in the given timeout.
+     *
+     * @param queryId The identifier of the query to check
+     * @param offset The starting chunk offset.
+     * @param limit The quantity of chunks relative to the offset to wait for
+     * @param timeout The duration to wait for the engine have results produced.
+     * @param allowLessThan Whether to return early when the available chunks is less than {@code offset + limit}
+     * @return The last {@link DataCloudQueryStatus} the server replied with.
+     */
+    public DataCloudQueryStatus waitForChunksAvailable(
+            String queryId, long offset, long limit, Duration timeout, boolean allowLessThan)
+            throws DataCloudJDBCException {
+        return executor.waitForChunksAvailable(queryId, offset, limit, timeout, allowLessThan);
+    }
+
     /**
      * Checks if all the query's results are ready, the row count and chunk count will be stable.
      * @param queryId The identifier of the query to check
      * @param timeout The duration to wait for the engine have results produced.
-     * @return The final {@link DataCloudQueryStatus} the server replied with.
+     * @return The last {@link DataCloudQueryStatus} the server replied with.
      */
     public DataCloudQueryStatus waitForResultsProduced(String queryId, Duration timeout) throws DataCloudJDBCException {
         return executor.waitForResultsProduced(queryId, timeout);

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudStatement.java

@@ -15,6 +15,8 @@
  */
 package com.salesforce.datacloud.jdbc.core;
 
+import static com.salesforce.datacloud.jdbc.core.HyperGrpcClientExecutor.HYPER_MAX_ROW_LIMIT_BYTE_SIZE;
+import static com.salesforce.datacloud.jdbc.core.HyperGrpcClientExecutor.HYPER_MIN_ROW_LIMIT_BYTE_SIZE;
 import static com.salesforce.datacloud.jdbc.util.PropertiesExtensions.getIntegerOrDefault;
 import static com.salesforce.datacloud.jdbc.util.PropertiesExtensions.optional;
 
@@ -32,6 +34,8 @@
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.time.Duration;
+import lombok.AccessLevel;
+import lombok.Getter;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
@@ -53,7 +57,9 @@ public class DataCloudStatement implements Statement, AutoCloseable {
 
     public DataCloudStatement(@NonNull DataCloudConnection connection) {
         this.dataCloudConnection = connection;
-        this.queryTimeout = getIntegerOrDefault(connection.getProperties(), QUERY_TIMEOUT, DEFAULT_QUERY_TIMEOUT);
+        val properties = connection.getProperties();
+        this.queryTimeout = getIntegerOrDefault(properties, QUERY_TIMEOUT, DEFAULT_QUERY_TIMEOUT);
+        this.targetMaxBytes = getIntegerOrDefault(properties, Constants.BYTE_LIMIT, HYPER_MAX_ROW_LIMIT_BYTE_SIZE);
     }
 
     protected QueryStatusListener listener;
@@ -66,6 +72,7 @@ protected boolean useSync() {
 
     private HyperGrpcClientExecutor getQueryExecutor() {
         return dataCloudConnection.getExecutor().toBuilder()
+                .byteLimit(targetMaxBytes)
                 .queryTimeout(getQueryTimeout())
                 .build();
     }
@@ -117,6 +124,7 @@ public DataCloudResultSet executeSyncQuery(String sql) throws SQLException {
         return executeSyncQuery(sql, client);
     }
 
+    @Deprecated
     protected DataCloudResultSet executeSyncQuery(String sql, HyperGrpcClientExecutor client) throws SQLException {
         listener = SyncQueryStatusListener.of(sql, client);
         resultSet = listener.generateResultSet();
@@ -133,7 +141,12 @@ public DataCloudResultSet executeAdaptiveQuery(String sql) throws SQLException {
 
     protected DataCloudResultSet executeAdaptiveQuery(String sql, HyperGrpcClientExecutor client, Duration timeout)
             throws SQLException {
-        listener = AdaptiveQueryStatusListener.of(sql, client, resolveQueryTimeout(timeout));
+        val queryTimeout = resolveQueryTimeout(timeout);
+
+        listener = targetMaxRows > 0
+                ? AdaptiveQueryStatusListener.of(sql, client, queryTimeout, targetMaxRows)
+                : AdaptiveQueryStatusListener.of(sql, client, queryTimeout);
+
         resultSet = listener.generateResultSet();
         log.info("executeAdaptiveQuery completed. queryId={}", listener.getQueryId());
         return (DataCloudResultSet) resultSet;
@@ -177,6 +190,54 @@ public int getMaxFieldSize() {
     @Override
     public void setMaxFieldSize(int max) {}
 
+    @Getter(AccessLevel.PACKAGE)
+    private int targetMaxRows;
+
+    @Getter(AccessLevel.PACKAGE)
+    private int targetMaxBytes;
+
+    public void clearResultSetConstraints() throws DataCloudJDBCException {
+        setResultSetConstraints(0, HYPER_MAX_ROW_LIMIT_BYTE_SIZE);
+    }
+
+    private void checkResultSetConstraints(int maxRows, int maxBytes) throws DataCloudJDBCException {
+        if (maxRows < 0) {
+            throw new DataCloudJDBCException(
+                    "setResultSetConstraints maxRows must be set to 0 to be disabled but was " + maxRows);
+        }
+
+        if (maxBytes < HYPER_MIN_ROW_LIMIT_BYTE_SIZE || maxBytes > HYPER_MAX_ROW_LIMIT_BYTE_SIZE) {
+            throw new DataCloudJDBCException(String.format(
+                    "The specified maxBytes (%d) must satisfy the following constraints: %d >= x >= %d",
+                    maxBytes, HYPER_MIN_ROW_LIMIT_BYTE_SIZE, HYPER_MAX_ROW_LIMIT_BYTE_SIZE));
+        }
+    }
+
+    /**
+     * Sets the constraints that would limit the number of rows and overall bytes in any ResultSet object generated by this Statement.
+     * This is used to tell the database the maximum number of rows and bytes to return, you may get less than expected because of this.
+     *
+     * @param maxRows The target maximum number of rows a ResultSet can have, zero means there is no limit.
+     * @param maxBytes The target maximum byte size a ResultSet can be,
+     *                 must fall in the range {@link HyperGrpcClientExecutor#HYPER_MIN_ROW_LIMIT_BYTE_SIZE}
+     *                 and {@link HyperGrpcClientExecutor#HYPER_MAX_ROW_LIMIT_BYTE_SIZE}
+     * @throws DataCloudJDBCException If the target maximum byte size is outside the aforementioned range
+     */
+    public void setResultSetConstraints(int maxRows, int maxBytes) throws DataCloudJDBCException {
+        checkResultSetConstraints(maxRows, maxBytes);
+
+        targetMaxRows = maxRows;
+        targetMaxBytes = maxBytes;
+    }
+
+    /**
+     * @see DataCloudStatement#setResultSetConstraints
+     * @param maxRows The target maximum number of rows a ResultSet can have, zero means there is no limit.
+     */
+    public void setResultSetConstraints(int maxRows) throws DataCloudJDBCException {
+        setResultSetConstraints(maxRows, HYPER_MAX_ROW_LIMIT_BYTE_SIZE);
+    }
+
     @Override
     public int getMaxRows() {
         return 0;