revamp query status polling to take advantage of getQueryInfo streaming

.github/workflows/build.yml

@@ -37,3 +37,13 @@ jobs:
             ${{ runner.os }}-hyper-${{ env.HYPER_VERSION }}
       - name: Maven package
         run: mvn $MAVEN_ARGS clean package --file pom.xml
+      - name: Upload hyper logs on failure
+        if: ${{ failure() || steps.main-process.outcome == 'failure' }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: test-results
+          path: |
+            **/target/surefire-reports/*.txt
+            **/target/surefire-reports/*.xml
+            target/hyper/*.log
+          retention-days: 5

.hooks/pre-commit

@@ -3,4 +3,8 @@ set -e
 
 echo '[git pre-commit] mvn spotless:apply sortpom:sort'
 MAVEN_OPTS='-Dorg.slf4j.simpleLogger.defaultLogLevel=error' mvn spotless:apply sortpom:sort
-git add --update
+
+staged_files=$(git diff --name-only --cached)
+for file in $staged_files; do
+    git add "$file"
+done

README.md

@@ -32,6 +32,15 @@ mvn clean install
 
 ## Usage
 
+> [!INFO]
+> Our API is versioned based on semantic versioning rules around our supported API.
+> This supported API includes:
+> 1. Any construct available through the JDBC specification we have implemented
+> 2. The DataCloudQueryStatus class
+> 3. The public methods in DataCloudConnection, DataCloudStatement, DataCloudResultSet, and DataCloudPreparedStatement -- note that these will be refactored to be interfaces that will make the API more obvious in the near future
+>
+> Usage of any other public classes or methods not listed above should be considered relatively unsafe, though we will strive to not make changes and will use semantic versioning from 1.0.0 and on.
+
 ### Connection string
 
 Use `jdbc:salesforce-datacloud://login.salesforce.com`

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/auth/AuthenticationSettings.java

@@ -25,7 +25,6 @@
 import com.salesforce.datacloud.jdbc.util.PropertiesExtensions;
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.sql.SQLException;
 import java.util.Properties;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -37,7 +36,7 @@
 
 @Getter
 public abstract class AuthenticationSettings {
-    public static AuthenticationSettings of(@NonNull Properties properties) throws SQLException {
+    public static AuthenticationSettings of(@NonNull Properties properties) throws DataCloudJDBCException {
         checkNotEmpty(properties);
         checkHasAllRequired(properties);
 
@@ -72,14 +71,14 @@ private static boolean hasRefreshToken(Properties properties) {
         return hasAll(properties, Keys.REFRESH_TOKEN_KEYS);
     }
 
-    private static void checkNotEmpty(@NonNull Properties properties) throws SQLException {
+    private static void checkNotEmpty(@NonNull Properties properties) throws DataCloudJDBCException {
         if (properties.isEmpty()) {
             throw new DataCloudJDBCException(
                     Messages.PROPERTIES_EMPTY, "28000", new IllegalArgumentException(Messages.PROPERTIES_EMPTY));
         }
     }
 
-    private static void checkHasAllRequired(Properties properties) throws SQLException {
+    private static void checkHasAllRequired(Properties properties) throws DataCloudJDBCException {
         if (hasAll(properties, Keys.REQUIRED_KEYS)) {
             return;
         }
@@ -91,15 +90,15 @@ private static void checkHasAllRequired(Properties properties) throws SQLExcepti
         throw new DataCloudJDBCException(missing, "28000", new IllegalArgumentException(missing));
     }
 
-    final URI getLoginUri() throws SQLException {
+    final URI getLoginUri() throws DataCloudJDBCException {
         try {
             return new URI(loginUrl);
         } catch (URISyntaxException ex) {
             throw new DataCloudJDBCException(ex.getMessage(), "28000", ex);
         }
     }
 
-    protected AuthenticationSettings(@NonNull Properties properties) throws SQLException {
+    protected AuthenticationSettings(@NonNull Properties properties) throws DataCloudJDBCException {
         checkNotEmpty(properties);
 
         this.relevantProperties = copy(properties, Keys.ALL);
@@ -172,7 +171,7 @@ protected static class Messages {
 
 @Getter
 class PasswordAuthenticationSettings extends AuthenticationSettings {
-    protected PasswordAuthenticationSettings(@NonNull Properties properties) throws SQLException {
+    protected PasswordAuthenticationSettings(@NonNull Properties properties) throws DataCloudJDBCException {
         super(properties);
 
         this.password = required(this.getRelevantProperties(), Keys.PASSWORD);
@@ -185,7 +184,7 @@ protected PasswordAuthenticationSettings(@NonNull Properties properties) throws
 
 @Getter
 class PrivateKeyAuthenticationSettings extends AuthenticationSettings {
-    protected PrivateKeyAuthenticationSettings(@NonNull Properties properties) throws SQLException {
+    protected PrivateKeyAuthenticationSettings(@NonNull Properties properties) throws DataCloudJDBCException {
         super(properties);
 
         this.privateKey = required(this.getRelevantProperties(), Keys.PRIVATE_KEY);
@@ -198,7 +197,7 @@ protected PrivateKeyAuthenticationSettings(@NonNull Properties properties) throw
 
 @Getter
 class RefreshTokenAuthenticationSettings extends AuthenticationSettings {
-    protected RefreshTokenAuthenticationSettings(@NonNull Properties properties) throws SQLException {
+    protected RefreshTokenAuthenticationSettings(@NonNull Properties properties) throws DataCloudJDBCException {
         super(properties);
 
         this.refreshToken = required(this.getRelevantProperties(), Keys.REFRESH_TOKEN);

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/auth/AuthenticationStrategy.java

@@ -30,12 +30,12 @@
 import org.apache.commons.lang3.StringUtils;
 
 interface AuthenticationStrategy {
-    static AuthenticationStrategy of(@NonNull Properties properties) throws SQLException {
+    static AuthenticationStrategy of(@NonNull Properties properties) throws DataCloudJDBCException {
         val settings = AuthenticationSettings.of(properties);
         return of(settings);
     }
 
-    static AuthenticationStrategy of(@NonNull AuthenticationSettings settings) throws SQLException {
+    static AuthenticationStrategy of(@NonNull AuthenticationSettings settings) throws DataCloudJDBCException {
         if (settings instanceof PasswordAuthenticationSettings) {
             return new PasswordAuthenticationStrategy((PasswordAuthenticationSettings) settings);
         } else if (settings instanceof PrivateKeyAuthenticationSettings) {
@@ -67,7 +67,7 @@ class Keys {
 }
 
 abstract class SharedAuthenticationStrategy implements AuthenticationStrategy {
-    protected final FormCommand.Builder builder(HttpCommandPath path) throws SQLException {
+    protected final FormCommand.Builder builder(HttpCommandPath path) throws DataCloudJDBCException {
         val settings = getSettings();
         val builder = FormCommand.builder();
 
@@ -95,7 +95,7 @@ class PasswordAuthenticationStrategy extends SharedAuthenticationStrategy {
      * password flow docs</a>
      */
     @Override
-    public FormCommand buildAuthenticate() throws SQLException {
+    public FormCommand buildAuthenticate() throws DataCloudJDBCException {
         val builder = super.builder(HttpCommandPath.AUTHENTICATE);
 
         builder.bodyEntry(Keys.GRANT_TYPE, GRANT_TYPE);
@@ -122,7 +122,7 @@ class RefreshTokenAuthenticationStrategy extends SharedAuthenticationStrategy {
      * token flow docs</a>
      */
     @Override
-    public FormCommand buildAuthenticate() throws SQLException {
+    public FormCommand buildAuthenticate() throws DataCloudJDBCException {
         val builder = super.builder(HttpCommandPath.AUTHENTICATE);
 
         builder.bodyEntry(Keys.GRANT_TYPE, GRANT_TYPE);
@@ -147,7 +147,7 @@ class PrivateKeyAuthenticationStrategy extends SharedAuthenticationStrategy {
      * docs</a>
      */
     @Override
-    public FormCommand buildAuthenticate() throws SQLException {
+    public FormCommand buildAuthenticate() throws DataCloudJDBCException {
         val builder = super.builder(HttpCommandPath.AUTHENTICATE);
 
         builder.bodyEntry(Keys.GRANT_TYPE, GRANT_TYPE);

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/auth/DataCloudTokenProcessor.java

@@ -160,7 +160,7 @@ private static AuthenticationResponseWithError throwExceptionOnError(
         return response;
     }
 
-    public static DataCloudTokenProcessor of(Properties properties) throws SQLException {
+    public static DataCloudTokenProcessor of(Properties properties) throws DataCloudJDBCException {
         val settings = AuthenticationSettings.of(properties);
         val strategy = AuthenticationStrategy.of(settings);
         val client = ClientBuilder.buildOkHttpClient(properties);

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/auth/PrivateKeyHelpers.java

@@ -63,7 +63,7 @@ public static Audience of(String url) throws SQLException {
 
 @UtilityClass
 class JwtParts {
-    public static String buildJwt(PrivateKeyAuthenticationSettings settings) throws SQLException {
+    public static String buildJwt(PrivateKeyAuthenticationSettings settings) throws DataCloudJDBCException {
         try {
             Instant now = Instant.now();
             Audience audience = Audience.of(settings.getLoginUrl());

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

@@ -15,10 +15,6 @@
  */
 package com.salesforce.datacloud.jdbc.core;
 
-import static com.salesforce.datacloud.jdbc.util.Constants.LOGIN_URL;
-import static com.salesforce.datacloud.jdbc.util.Constants.USER;
-import static com.salesforce.datacloud.jdbc.util.Constants.USER_NAME;
-
 import com.salesforce.datacloud.jdbc.auth.AuthenticationSettings;
 import com.salesforce.datacloud.jdbc.auth.DataCloudTokenProcessor;
 import com.salesforce.datacloud.jdbc.auth.TokenProcessor;
@@ -30,10 +26,17 @@
 import com.salesforce.datacloud.jdbc.interceptor.DataspaceHeaderInterceptor;
 import com.salesforce.datacloud.jdbc.interceptor.HyperExternalClientContextHeaderInterceptor;
 import com.salesforce.datacloud.jdbc.interceptor.HyperWorkloadHeaderInterceptor;
-import com.salesforce.datacloud.jdbc.interceptor.TracingHeadersInterceptor;
 import com.salesforce.datacloud.jdbc.util.Unstable;
+import com.salesforce.datacloud.query.v3.DataCloudQueryStatus;
 import io.grpc.ClientInterceptor;
 import io.grpc.ManagedChannelBuilder;
+import lombok.AccessLevel;
+import lombok.Builder;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+
 import java.sql.Array;
 import java.sql.Blob;
 import java.sql.CallableStatement;
@@ -49,7 +52,7 @@
 import java.sql.Savepoint;
 import java.sql.Statement;
 import java.sql.Struct;
-import java.util.ArrayList;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -59,13 +62,10 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import lombok.AccessLevel;
-import lombok.Builder;
-import lombok.Getter;
-import lombok.NonNull;
-import lombok.Setter;
-import lombok.extern.slf4j.Slf4j;
-import lombok.val;
+
+import static com.salesforce.datacloud.jdbc.util.Constants.LOGIN_URL;
+import static com.salesforce.datacloud.jdbc.util.Constants.USER;
+import static com.salesforce.datacloud.jdbc.util.Constants.USER_NAME;
 
 @Slf4j
 @Builder(access = AccessLevel.PACKAGE)
@@ -82,11 +82,7 @@ public class DataCloudConnection implements Connection, AutoCloseable {
     @NonNull @Builder.Default
     private final Properties properties = new Properties();
 
-    @Getter(AccessLevel.PACKAGE)
-    @Setter
-    @Builder.Default
-    private List<ClientInterceptor> interceptors = new ArrayList<>();
-
+    @Unstable
     @Getter(AccessLevel.PACKAGE)
     @NonNull private final HyperGrpcClientExecutor executor;
 
@@ -106,33 +102,12 @@ public static DataCloudConnection fromChannel(@NonNull ManagedChannelBuilder<?>
                 .build();
     }
 
-    /** This flow is not supported by the JDBC Driver Manager, only use it if you know what you're doing. */
-    public static DataCloudConnection fromTokenSupplier(
-            AuthorizationHeaderInterceptor authInterceptor, @NonNull String host, int port, Properties properties)
-            throws SQLException {
-        val channel = ManagedChannelBuilder.forAddress(host, port);
-        return fromTokenSupplier(authInterceptor, channel, properties);
-    }
-
-    /** This flow is not supported by the JDBC Driver Manager, only use it if you know what you're doing. */
-    public static DataCloudConnection fromTokenSupplier(
-            AuthorizationHeaderInterceptor authInterceptor, ManagedChannelBuilder<?> builder, Properties properties)
-            throws SQLException {
-        val interceptors = getClientInterceptors(authInterceptor, properties);
-        val executor = HyperGrpcClientExecutor.of(builder.intercept(interceptors), properties);
-
-        return DataCloudConnection.builder()
-                .executor(executor)
-                .properties(properties)
-                .build();
-    }
-
     /**
      * Initializes a list of interceptors that handle channel level concerns that can be defined through properties
      * @param properties - The connection properties
      * @return a list of client interceptors
      */
-    static List<ClientInterceptor> getPropertyDerivedClientInterceptors(Properties properties) {
+    private static List<ClientInterceptor> getPropertyDerivedClientInterceptors(Properties properties) {
         return Stream.of(
                         HyperExternalClientContextHeaderInterceptor.of(properties),
                         HyperWorkloadHeaderInterceptor.of(properties),
@@ -141,22 +116,13 @@ static List<ClientInterceptor> getPropertyDerivedClientInterceptors(Properties p
                 .collect(Collectors.toList());
     }
 
-    /**
-     * Initializes the full set of client interceptors from property handling to tracing and auth
-     * @param authInterceptor an optional auth interceptor, is allowed to be null
-     * @param properties the connection properties
-     * @return a list of client interceptors
-     */
-    static List<ClientInterceptor> getClientInterceptors(
-            AuthorizationHeaderInterceptor authInterceptor, Properties properties) {
-        val list = getPropertyDerivedClientInterceptors(properties);
-        list.add(0, TracingHeadersInterceptor.of());
-        if (authInterceptor != null) {
-            list.add(0, authInterceptor);
+    private static DataCloudTokenProcessor getDataCloudTokenProcessor(Properties properties)
+            throws DataCloudJDBCException {
+        if (!AuthenticationSettings.hasAny(properties)) {
+            throw new DataCloudJDBCException("No authentication settings provided");
         }
-        ;
-        log.info("Registering interceptor. interceptor={}", list);
-        return list;
+
+        return DataCloudTokenProcessor.of(properties);
     }
 
     public static DataCloudConnection of(String url, Properties properties) throws SQLException {
@@ -165,17 +131,15 @@ public static DataCloudConnection of(String url, Properties properties) throws S
         connectionString.withParameters(properties);
         properties.setProperty(LOGIN_URL, connectionString.getLoginUrl());
 
-        if (!AuthenticationSettings.hasAny(properties)) {
-            throw new DataCloudJDBCException("No authentication settings provided");
-        }
-
-        val tokenProcessor = DataCloudTokenProcessor.of(properties);
+        val tokenProcessor = getDataCloudTokenProcessor(properties);
+        val authInterceptor = AuthorizationHeaderInterceptor.of(tokenProcessor);
 
         val host = tokenProcessor.getDataCloudToken().getTenantUrl();
         val builder = ManagedChannelBuilder.forAddress(host, DEFAULT_PORT);
-        val authInterceptor = AuthorizationHeaderInterceptor.of(tokenProcessor);
 
-        val interceptors = getClientInterceptors(authInterceptor, properties);
+        val interceptors = getPropertyDerivedClientInterceptors(properties);
+        interceptors.add(0, authInterceptor);
+
         val executor = HyperGrpcClientExecutor.of(builder.intercept(interceptors), properties);
 
         return DataCloudConnection.builder()
@@ -221,22 +185,44 @@ public DataCloudResultSet getRowBasedResultSet(String queryId, long offset, long
         return StreamingResultSet.of(queryId, executor, iterator);
     }
 
-    @Unstable
     public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId, long limit) {
         log.info("Get chunk-based result set. queryId={}, chunkId={}, limit={}", queryId, chunkId, limit);
         val iterator = ChunkBased.of(executor, queryId, chunkId, limit);
         return StreamingResultSet.of(queryId, executor, iterator);
     }
 
-    @Unstable
     public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId) {
         return getChunkBasedResultSet(queryId, chunkId, 1);
     }
 
+    /**
+     * 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 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 or not to return early when the available rows is less than {@code offset + limit}
+     * @return The final {@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 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.
+     */
+    public DataCloudQueryStatus waitForResultsProduced(String queryId, Duration timeout) throws DataCloudJDBCException {
+        return executor.waitForResultsProduced(queryId, timeout);
+    }
+
     /**
      * Use this to determine when a given query is complete by filtering the responses and a subsequent findFirst()
      */
-    @Unstable
     public Stream<DataCloudQueryStatus> getQueryStatus(String queryId) {
         return executor.getQueryStatus(queryId);
     }