fix tests

Author: jschneidereit

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

@@ -184,13 +184,16 @@ 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);
     }
 
+    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
@@ -200,7 +203,6 @@ public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId, l
      * @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.
      */
-    @Unstable
     public DataCloudQueryStatus waitForRowsAvailable(
             String queryId, long offset, long limit, Duration timeout, boolean allowLessThan)
             throws DataCloudJDBCException {
@@ -213,20 +215,13 @@ public DataCloudQueryStatus waitForRowsAvailable(
      * @param timeout The duration to wait for the engine have results produced.
      * @return The final {@link DataCloudQueryStatus} the server replied with.
      */
-    @Unstable
     public DataCloudQueryStatus waitForResultsProduced(String queryId, Duration timeout) throws DataCloudJDBCException {
         return executor.waitForResultsProduced(queryId, timeout);
     }
 
-    @Unstable
-    public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId) {
-        return getChunkBasedResultSet(queryId, chunkId, 1);
-    }
-
     /**
      * 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);
     }

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

@@ -261,5 +261,15 @@ public void close() throws Exception {
         }
 
         channel.shutdown();
+
+        try {
+            channel.awaitTermination(5, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            log.error("Failed to shutdown channel within 5 seconds", e);
+        } finally {
+            if (!channel.isTerminated()) {
+                channel.shutdownNow();
+            }
+        }
     }
 }

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/partial/DataCloudQueryPolling.java

@@ -19,11 +19,6 @@
 import com.salesforce.datacloud.jdbc.util.Unstable;
 import com.salesforce.datacloud.query.v3.DataCloudQueryStatus;
 import io.grpc.StatusRuntimeException;
-import java.time.Duration;
-import java.time.Instant;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Predicate;
 import lombok.SneakyThrows;
 import lombok.experimental.UtilityClass;
 import lombok.extern.slf4j.Slf4j;
@@ -34,6 +29,12 @@
 import salesforce.cdp.hyperdb.v1.HyperServiceGrpc;
 import salesforce.cdp.hyperdb.v1.QueryInfoParam;
 
+import java.time.Duration;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Predicate;
+
 @Unstable
 @Slf4j
 @UtilityClass
@@ -43,16 +44,31 @@ public static DataCloudQueryStatus waitForRowsAvailable(
             String queryId,
             long offset,
             long limit,
-            Duration timeout,
+            Duration timeoutDuration,
             boolean allowLessThan)
             throws DataCloudJDBCException {
-        return waitForQueryStatus(stub, queryId, timeout, status -> {
+
+        Predicate<DataCloudQueryStatus> predicate = status -> {
             if (allowLessThan) {
                 return status.getRowCount() > offset;
             } else {
                 return status.getRowCount() >= offset + limit;
             }
-        });
+        };
+
+        val result = waitForQueryStatus(stub, queryId, timeoutDuration, predicate);
+
+        if (predicate.test(result)) {
+            return result;
+        } else {
+            if (allowLessThan) {
+                throw new DataCloudJDBCException(
+                        "Timed out waiting for new rows to be available. queryId=" + queryId + ", status=" + result);
+            } else {
+                throw new DataCloudJDBCException(
+                        "Timed out waiting for enough rows to be available. queryId=" + queryId + ", status=" + result);
+            }
+        }
     }
 
     public static DataCloudQueryStatus waitForResultsProduced(
@@ -64,17 +80,21 @@ public static DataCloudQueryStatus waitForResultsProduced(
     public static DataCloudQueryStatus waitForQueryStatus(
             HyperServiceGrpc.HyperServiceBlockingStub stub,
             String queryId,
-            Duration timeout,
+            Duration timeoutDuration,
             Predicate<DataCloudQueryStatus> predicate)
             throws DataCloudJDBCException {
         val last = new AtomicReference<DataCloudQueryStatus>();
-        val deadline = Instant.now().plus(timeout);
-        val attempts = new AtomicInteger(1);
+        val deadline = Instant.now().plus(timeoutDuration);
+        val attempts = new AtomicInteger(0);
 
-        val retry = new RetryPolicy<DataCloudQueryStatus>()
-                .withMaxDuration(timeout)
-                .handle(StatusRuntimeException.class)
+        val retryPolicy = new RetryPolicy<DataCloudQueryStatus>()
+                .withMaxDuration(timeoutDuration)
                 .handleIf(e -> {
+                    if (!(e instanceof StatusRuntimeException)) {
+                        log.error("Got an unexpected exception when getting query status for queryId={}", queryId, e);
+                        return false;
+                    }
+
                     if (last.get() == null) {
                         log.error(
                                 "Failed to get query status response, will not try again. queryId={}, attempts={}",
@@ -103,7 +123,7 @@ public static DataCloudQueryStatus waitForQueryStatus(
                 });
 
         try {
-            return Failsafe.with(retry)
+            return Failsafe.with(retryPolicy)
                     .get(() -> waitForQueryStatusWithoutRetry(stub, queryId, deadline, last, attempts, predicate));
         } catch (FailsafeException ex) {
             throw new DataCloudJDBCException(
@@ -126,14 +146,21 @@ static DataCloudQueryStatus waitForQueryStatusWithoutRetry(
         times.getAndIncrement();
         val param = QueryInfoParam.newBuilder().setQueryId(queryId).build();
         while (Instant.now().isBefore(deadline)) {
-            val info = stub.getQueryInfo(param);
+            val timeRemaining = Duration.between(Instant.now(), deadline);
+            val info = stub.withDeadlineAfter(timeRemaining).getQueryInfo(param);
             while (info.hasNext()) {
-                val next = DataCloudQueryStatus.of(info.next());
-                next.ifPresent(last::set);
-                if (predicate.test(last.get())) {
+                val matched = DataCloudQueryStatus.of(info.next())
+                        .map(next -> {
+                            last.set(next);
+                            return predicate.test(next);
+                        })
+                        .orElse(false);
+
+                if (matched) {
                     return last.get();
                 }
             }
+            log.info("end of info stream, maybe starting a new one: {}", last.get());
         }
         return last.get();
     }

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/interceptor/HeaderMutatingClientInterceptor.java

@@ -27,7 +27,7 @@
 
 @FunctionalInterface
 public interface HeaderMutatingClientInterceptor extends ClientInterceptor {
-    void mutate(final Metadata headers);
+    void mutate(final Metadata headers) throws DataCloudJDBCException;
 
     @Override
     default <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(

jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/DataCloudStatementFunctionalTest.java

@@ -36,9 +36,10 @@ public class DataCloudStatementFunctionalTest extends HyperTestBase {
     @SneakyThrows
     public void canCancelStatementQuery() {
         try (val server = configWithSleep.start();
-                val statement = server.getConnection().createStatement().unwrap(DataCloudStatement.class)) {
+             val statement = server.getConnection().createStatement().unwrap(DataCloudStatement.class);
+             val client = server.getRawClient()) {
             statement.execute("select pg_sleep(5000000);");
-            val client = server.getRawClient();
+
             val queryId = statement.getQueryId();
             val a = client.getQueryStatus(queryId).findFirst().get();
             assertThat(a.getCompletionStatus()).isEqualTo(DataCloudQueryStatus.CompletionStatus.RUNNING);
@@ -56,10 +57,12 @@ public void canCancelPreparedStatementQuery() {
         try (val server = configWithSleep.start();
                 val statement = server.getConnection()
                         .prepareStatement("select pg_sleep(?)")
-                        .unwrap(DataCloudPreparedStatement.class)) {
+                        .unwrap(DataCloudPreparedStatement.class);
+             val client = server.getRawClient()) {
+
             statement.setInt(1, 5000000);
             statement.execute();
-            val client = server.getRawClient();
+
             val queryId = statement.getQueryId();
             val a = client.getQueryStatus(queryId).findFirst().get();
             assertThat(a.getCompletionStatus()).isEqualTo(DataCloudQueryStatus.CompletionStatus.RUNNING);
@@ -75,14 +78,13 @@ public void canCancelPreparedStatementQuery() {
     @SneakyThrows
     public void canCancelAnotherQueryById() {
         try (val server = configWithSleep.start();
-                val statement = server.getConnection().createStatement().unwrap(DataCloudStatement.class);
-                val cancel = server.getConnection().unwrap(DataCloudConnection.class)) {
+             val statement = server.getConnection().createStatement().unwrap(DataCloudStatement.class);
+             val cancel = server.getConnection().unwrap(DataCloudConnection.class);
+             val client = server.getRawClient()) {
 
             statement.execute("select pg_sleep(5000000);");
             val queryId = statement.getQueryId();
 
-            val client = server.getRawClient();
-
             val a = client.getQueryStatus(queryId).findFirst().get();
             assertThat(a.getCompletionStatus()).isEqualTo(DataCloudQueryStatus.CompletionStatus.RUNNING);
 
@@ -111,6 +113,8 @@ public void forwardAndReadOnly() {
 
             assertThat(rs.getType()).isEqualTo(ResultSet.TYPE_FORWARD_ONLY);
             assertThat(rs.getConcurrency()).isEqualTo(ResultSet.CONCUR_READ_ONLY);
+
+            assertThat(rs.getRow()).isEqualTo(0);
         });
     }
 

jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/StreamingResultSetTest.java

@@ -23,6 +23,7 @@
 import com.salesforce.datacloud.jdbc.util.ThrowingBiFunction;
 import io.grpc.StatusRuntimeException;
 import java.sql.SQLException;
+import java.time.Duration;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 import lombok.SneakyThrows;
@@ -59,15 +60,17 @@ public void exercisePreparedStatement() {
         val expected = new AtomicInteger(0);
 
         assertWithConnection(conn -> {
-            try (val statement = conn.prepareStatement(sql)) {
+            try (val statement = conn.prepareStatement(sql).unwrap(DataCloudPreparedStatement.class)) {
                 statement.setInt(1, large);
 
                 val rs = statement.executeQuery();
+                conn.waitForResultsProduced(statement.getQueryId(), Duration.ofSeconds(30));
                 assertThat(rs).isInstanceOf(StreamingResultSet.class);
                 assertThat(((StreamingResultSet) rs).isReady()).isTrue();
 
                 while (rs.next()) {
                     assertEachRowIsTheSame(rs, expected);
+                    assertThat(rs.getRow()).isEqualTo(expected.get());
                 }
             }
         });
@@ -85,6 +88,7 @@ public void exerciseQueryMode(
 
         assertWithStatement(statement -> {
             val rs = queryMode.apply(statement, sql);
+
             assertThat(rs).isInstanceOf(StreamingResultSet.class);
             assertThat(rs.isReady()).isTrue();
 
@@ -133,23 +137,13 @@ private static Arguments deferred(
                     impl.apply(s, x);
 
                     if (wait) {
-                        waitUntilReady(s);
+                        val conn = s.getConnection().unwrap(DataCloudConnection.class);
+                        conn.waitForResultsProduced(s.getQueryId(), Duration.ofSeconds(30));
                     }
 
                     return (DataCloudResultSet) s.getResultSet();
                 };
         return arguments(named(String.format("%s; getResultSet -> DataCloudResultSet", name), deferred), size);
     }
 
-    @SneakyThrows
-    static boolean waitUntilReady(DataCloudStatement statement) {
-        while (!statement.isReady()) {
-            try {
-                Thread.sleep(1000);
-            } catch (InterruptedException e) {
-                return false;
-            }
-        }
-        return true;
-    }
 }

jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/partial/DataCloudQueryPollingMockTests.java

@@ -0,0 +1,50 @@
+/*
+ * 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.core.partial;
+
+import com.salesforce.datacloud.jdbc.core.DataCloudStatement;
+import com.salesforce.datacloud.jdbc.core.HyperGrpcTestBase;
+import lombok.SneakyThrows;
+import lombok.val;
+import org.assertj.core.api.AssertionsForClassTypes;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import salesforce.cdp.hyperdb.v1.HyperServiceGrpc;
+
+import java.time.Duration;
+
+import static org.grpcmock.GrpcMock.calledMethod;
+import static org.grpcmock.GrpcMock.times;
+import static org.grpcmock.GrpcMock.verifyThat;
+
+public class DataCloudQueryPollingMockTests extends HyperGrpcTestBase {
+    @Test
+    @SneakyThrows
+    @Disabled
+    void getQueryInfoDoesNotRetryIfFailureToConnect() {
+        try (val connection = getInterceptedClientConnection();
+                val statement = connection.createStatement().unwrap(DataCloudStatement.class)) {
+            statement.execute("select * from nonsense");
+
+            verifyThat(calledMethod(HyperServiceGrpc.getGetQueryInfoMethod()), times(0));
+
+            AssertionsForClassTypes.assertThatThrownBy(
+                    () -> connection.waitForResultsProduced(statement.getQueryId(), Duration.ofSeconds(30)));
+
+            verifyThat(calledMethod(HyperServiceGrpc.getGetQueryInfoMethod()), times(1));
+        }
+    }
+}