Adjust fromChannel method to minimize injected headers and add examples

mkaufmann · mkaufmann · commit 1555d4ee5129 · 2025-03-11T20:33:11.000+01:00
The intent of this change was to analyze how to use the JDBC driver in a least invasive way (compared to manual gRPC calling).
To reach that minimum the `DataCloudConnection.fromChannel` method was adjusted to only add the must have interceptors, which - without usage of the corresponding connection parameters - are complete no ops.

To demonstrate this flow a new `examples` test package was create that demonstrates the approach and serves as documentation.
To write that test case (and minimize the use of helpers which would hide complexity) the test hyper instance variable with its getPort method was exposed as public.

Drive-By:t
- The hyper test instance was made non-static to avoid problems in parallel test execution.
- A sample implementation of using the row based API to serve row based paginated results was also added and commented. With the parallel ongoing work we should see simplification of it.
diff --git a/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudConnection.java b/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudConnection.java
@@ -90,9 +90,14 @@ public class DataCloudConnection implements Connection, AutoCloseable {
     @Getter(AccessLevel.PACKAGE)
     @NonNull private final HyperGrpcClientExecutor executor;
 
+    /**
+     * This creates a Data Cloud connection with minimal adjustments to the channels.
+     * The only added interceptors are those for handling connection parameters that influence headers.
+     * This will not provide auth / tracing, users of this API are expected to wire their own
+     */
     public static DataCloudConnection fromChannel(@NonNull ManagedChannelBuilder<?> builder, Properties properties)
             throws SQLException {
-        val interceptors = getClientInterceptors(null, properties);
+        val interceptors = getPropertyDerivedClientInterceptors(properties);
         val executor = HyperGrpcClientExecutor.of(builder.intercept(interceptors), properties);
 
         return DataCloudConnection.builder()
@@ -122,19 +127,38 @@ public static DataCloudConnection fromTokenSupplier(
                 .build();
     }
 
-    static List<ClientInterceptor> getClientInterceptors(
-            AuthorizationHeaderInterceptor authInterceptor, Properties properties) {
+    /**
+     * 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) {
         return Stream.of(
-                        authInterceptor,
-                        TracingHeadersInterceptor.of(),
                         HyperExternalClientContextHeaderInterceptor.of(properties),
                         HyperWorkloadHeaderInterceptor.of(properties),
                         DataspaceHeaderInterceptor.of(properties))
                 .filter(Objects::nonNull)
-                .peek(t -> log.info("Registering interceptor. interceptor={}", t))
                 .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(TracingHeadersInterceptor.of());
+        if (authInterceptor != null) {
+            list.add(authInterceptor);
+        }
+        ;
+        log.info("Registering interceptor. interceptor={}", list);
+        return list;
+    }
+
     public static DataCloudConnection of(String url, Properties properties) throws SQLException {
         val connectionString = DataCloudConnectionString.of(url);
         addClientUsernameIfRequired(properties);
@@ -179,6 +203,7 @@ private DataCloudPreparedStatement getQueryPreparedStatement(String sql) {
     /**
      * Retrieves a collection of rows for the specified query once it is ready.
      * Use {@link #getQueryStatus(String)} to check if the query has produced results or finished execution before calling this method.
+     * You can get the Query Id from the executeQuery `DataCloudResultSet`.
      * <p>
      * When using {@link RowBased.Mode#FULL_RANGE}, this method does not handle pagination near the end of available rows.
      * The caller is responsible for calculating the correct offset and limit to avoid out-of-range errors.
diff --git a/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/interceptor/SingleHeaderMutatingClientInterceptor.java b/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/interceptor/SingleHeaderMutatingClientInterceptor.java
@@ -18,7 +18,7 @@
 import io.grpc.Metadata;
 import lombok.NonNull;
 
-interface SingleHeaderMutatingClientInterceptor extends HeaderMutatingClientInterceptor {
+public interface SingleHeaderMutatingClientInterceptor extends HeaderMutatingClientInterceptor {
     @NonNull Metadata.Key<String> getKey();
 
     @NonNull String getValue();
diff --git a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/AsyncStreamingResultSetTest.java b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/AsyncStreamingResultSetTest.java
@@ -36,7 +36,7 @@ public class AsyncStreamingResultSetTest extends HyperTestBase {
     @SneakyThrows
     public void testThrowsOnNonsenseQueryAsync() {
         val ex = Assertions.assertThrows(DataCloudJDBCException.class, () -> {
-            try (val connection = HyperTestBase.getHyperQueryConnection();
+            try (val connection = getHyperQueryConnection();
                     val statement = connection.createStatement().unwrap(DataCloudStatement.class)) {
                 val rs = statement.executeAsyncQuery("select * from nonsense");
                 waitUntilReady(statement);
diff --git a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/examples/SubmitQueryAndConsumeResultsTest.java b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/examples/SubmitQueryAndConsumeResultsTest.java
@@ -0,0 +1,181 @@
+/*
+ * Copyright (c) 2024, Salesforce, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.salesforce.datacloud.jdbc.examples;
+
+import static java.lang.Math.min;
+
+import com.salesforce.datacloud.jdbc.core.DataCloudConnection;
+import com.salesforce.datacloud.jdbc.core.DataCloudQueryStatus;
+import com.salesforce.datacloud.jdbc.core.DataCloudResultSet;
+import com.salesforce.datacloud.jdbc.core.partial.RowBased;
+import com.salesforce.datacloud.jdbc.hyper.HyperTestBase;
+import io.grpc.ManagedChannelBuilder;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.*;
+import lombok.extern.slf4j.Slf4j;
+import org.junit.jupiter.api.Test;
+
+/**
+ * This example uses a locally spawned Hyper instance to demonstrate best practices around connecting to Hyper.
+ * This consciously only uses the JDBC API in the core and no helpers (outside of this class) to provide self contained
+ * examples.
+ */
+@Slf4j
+public class SubmitQueryAndConsumeResultsTest extends HyperTestBase {
+    /**
+     * This example shows how to create a Data Cloud Connection while still having full control over concerns like
+     * authorization and tracing.
+     */
+    @Test
+    public void testBareBonesExecuteQuery() throws SQLException {
+        // The connection properties
+        Properties properties = new Properties();
+
+        // You can bring your own gRPC channels, setup in the way you like (mTLS / Plaintext / ...) and your own
+        // interceptors as well as executors.
+        ManagedChannelBuilder<?> channel = ManagedChannelBuilder.forAddress("127.0.0.1", instance.getPort())
+                .usePlaintext();
+
+        // Use the JDBC Driver interface
+        try (DataCloudConnection conn = DataCloudConnection.fromChannel(channel, properties)) {
+            try (Statement stmt = conn.createStatement()) {
+                ResultSet rs = stmt.executeQuery("SELECT s FROM generate_series(1,10) s");
+                while (rs.next()) {
+                    System.out.println("Retrieved value:" + rs.getLong(1));
+                }
+            }
+        }
+    }
+
+    /**
+     * Analyze the query status, as we have a query status we know that the query was last observed in a non failing
+     * state.
+     *
+     * Offset must always be larger or equal to get row count (which would happen for typical next based pagination)
+     */
+    private static long rowBasedStatusObjectRowsCheck(DataCloudQueryStatus queryStatus, long offset, long pageLimit) {
+        // Check if we can at least return some data
+        if (queryStatus.getRowCount() > offset) {
+            return min(queryStatus.getRowCount() - offset, pageLimit);
+        }
+        // A negative count signals that no data is available
+        return -1;
+    }
+
+    /**
+     * Checks if the query status signals that all results are produced
+     */
+    private static boolean allResultsProduced(DataCloudQueryStatus queryStatus) {
+        return queryStatus.isResultProduced() || queryStatus.isExecutionFinished();
+    }
+
+    /**
+     * This example shows how to use the row based pagination mode to get results segmented by approximate row count.
+     * For the example we access the results in 2 row ranges and have an implementation where the application doesn't
+     * know how many results would be produced in the end
+     */
+    @Test
+    public void testRowBasedPagination() throws SQLException {
+        final int pageRowLimit = 2;
+        long offset = 0;
+        long page = 0;
+
+        // The connection properties
+        Properties properties = new Properties();
+
+        // You can bring your own gRPC channels, setup in the way you like (mTLS / Plaintext / ...) and your own
+        // interceptors as well as executors.
+        ManagedChannelBuilder<?> channel = ManagedChannelBuilder.forAddress("127.0.0.1", instance.getPort())
+                .usePlaintext();
+
+        try (DataCloudConnection conn = DataCloudConnection.fromChannel(channel, properties)) {
+            // Submit the query and consume the initial page
+            String queryId;
+            try (Statement stmt = conn.createStatement()) {
+                log.warn("Executing query using a single `ExecuteQuery` RPC Call");
+                ResultSet rs = stmt.executeQuery("SELECT s FROM generate_series(1,11) s");
+                queryId = ((DataCloudResultSet) rs).getQueryId();
+                // For this result set we as a consumer must currently implement the pagination limit ourselves
+                int i = 0;
+                while (rs.next() && (i++ < pageRowLimit)) {
+                    ++offset;
+                    System.out.println("Retrieved value: " + rs.getLong(1) + " on page " + page);
+                }
+                ++page;
+            }
+
+            // Consume further pages until the full result is consumed (could also be done on a new connection if
+            // needed)
+            // NIT: We should provide an API on the original result set to access the `DataCloudQueryStatus` that way,
+            // if the query is already finished we don't need to do another network round-trip.
+            Optional<DataCloudQueryStatus> cachedStatus = Optional.empty();
+            while (true) {
+                // Try to make sure we have a status object
+                if (!cachedStatus.isPresent()) {
+                    // Identify if there is more data?
+                    long lambdaOffset = offset;
+                    // In case of query error this could throw an runtime exception
+                    // NIT: What is the timeout enforced here?
+                    log.warn("Fetching query status using a single `GetQueryInfo` RPC call");
+                    // NIT: Semantically I would want takeWhile here which is only available in Java 11
+                    cachedStatus = conn.getQueryStatus(queryId)
+                            .filter(queryStatus ->
+                                    (rowBasedStatusObjectRowsCheck(queryStatus, lambdaOffset, pageRowLimit) > 0)
+                                            || allResultsProduced(queryStatus))
+                            .findFirst();
+
+                    // Query is still running
+                    // NIT: Check how we should handle this in the presence of timeouts
+                    if (!cachedStatus.isPresent()) {
+                        continue;
+                    }
+                }
+
+                long availableRows = rowBasedStatusObjectRowsCheck(cachedStatus.get(), offset, pageRowLimit);
+                // Check if query completed and thus we can't produce more results
+                if (availableRows <= 0) {
+                    if (allResultsProduced(cachedStatus.get())) {
+                        break;
+                    } else {
+                        // NIT: Enable non sleep based flow
+                        try {
+                            Thread.sleep(100);
+                        } catch (InterruptedException e) {
+                            throw new RuntimeException(e);
+                        }
+                        // We need to fetch a new status in the next iteration
+                        cachedStatus = Optional.empty();
+                        continue;
+                    }
+                }
+
+                // At this point we know that rows are available
+                log.warn("Fetching query status using a single `GetQueryResult` RPC call");
+                try (ResultSet rs =
+                        conn.getRowBasedResultSet(queryId, offset, pageRowLimit, RowBased.Mode.SINGLE_RPC)) {
+                    while (rs.next()) {
+                        ++offset;
+                        System.out.println("Retrieved value: " + rs.getLong(1) + " on page " + page);
+                    }
+                    ++page;
+                }
+            }
+        }
+        log.warn("Completed");
+    }
+}
diff --git a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/hyper/HyperServerProcess.java b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/hyper/HyperServerProcess.java
@@ -97,7 +97,7 @@ public HyperServerProcess(HyperServerConfig.HyperServerConfigBuilder config) {
         }
     }
 
-    int getPort() {
+    public int getPort() {
         return port;
     }
 
diff --git a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/hyper/HyperTestBase.java b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/hyper/HyperTestBase.java
@@ -49,18 +49,18 @@
 @Slf4j
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public class HyperTestBase {
-    private static HyperServerProcess instance;
+    public HyperServerProcess instance;
 
     @SneakyThrows
-    public static void assertEachRowIsTheSame(ResultSet rs, AtomicInteger prev) {
+    public final void assertEachRowIsTheSame(ResultSet rs, AtomicInteger prev) {
         val expected = prev.incrementAndGet();
         val a = rs.getBigDecimal(1).intValue();
         assertThat(expected).isEqualTo(a);
     }
 
     @SafeVarargs
     @SneakyThrows
-    public static void assertWithConnection(
+    public final void assertWithConnection(
             ThrowingConsumer<DataCloudConnection> assertion, Map.Entry<String, String>... settings) {
         try (val connection =
                 getHyperQueryConnection(settings == null ? ImmutableMap.of() : ImmutableMap.ofEntries(settings))) {
@@ -70,7 +70,7 @@ public static void assertWithConnection(
 
     @SafeVarargs
     @SneakyThrows
-    public static void assertWithStatement(
+    public final void assertWithStatement(
             ThrowingConsumer<DataCloudStatement> assertion, Map.Entry<String, String>... settings) {
         try (val connection = getHyperQueryConnection(
                         settings == null ? ImmutableMap.of() : ImmutableMap.ofEntries(settings));
@@ -79,11 +79,11 @@ public static void assertWithStatement(
         }
     }
 
-    public static DataCloudConnection getHyperQueryConnection() {
+    public DataCloudConnection getHyperQueryConnection() {
         return getHyperQueryConnection(ImmutableMap.of());
     }
 
-    public static DataCloudConnection getHyperQueryConnection(Map<String, String> connectionSettings) {
+    public DataCloudConnection getHyperQueryConnection(Map<String, String> connectionSettings) {
         return instance.getConnection(connectionSettings);
     }
 
diff --git a/pom.xml b/pom.xml
@@ -19,7 +19,7 @@
         <flatten-maven-plugin.version>1.6.0</flatten-maven-plugin.version>
         <git-build-hook-maven-plugin.version>3.5.0</git-build-hook-maven-plugin.version>
         <!-- https://tableau.github.io/hyper-db/docs/releases#download -->
-        <hyperapi.version>0.0.21200.re11c8cb9</hyperapi.version>
+        <hyperapi.version>0.0.21408.rf5a406c0</hyperapi.version>
         <hyperd.directory>${project.build.directory}/hyper</hyperd.directory>
         <junit-bom.version>5.11.3</junit-bom.version>
         <maven-depndency-plugin.version>3.8.1</maven-depndency-plugin.version>

Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ public HyperServerProcess(HyperServerConfig.HyperServerConfigBuilder config) {`
`97`	`97`	`}`
`98`	`98`	`}`
`99`	`99`
`100`		`- int getPort() {`
	`100`	`+ public int getPort() {`
`101`	`101`	`return port;`
`102`	`102`	`}`
`103`	`103`