[PECOBLR-1131] Fix incorrect refetching of expired CloudFetch links when using Thrift protocol. (#1066)

## Description
[Thrift
protocol](https://github.com/databricks-eng/runtime/blob/master/sql/hive-thriftserver/if/TCLIService.thrift#L2186)
has a orientation field with values FETCH_NEXT, FETCH_PRIOR or
FETCH_FIRST. This field is always set to FETCH_NEXT resulting in
incorrect refetch. To fetch from a particular chunk index the Thrift
protocol requires the start row offset to be set. The chunk index and
start row offset information is available from the expired links. Use
the start row offset to fetch the links in the Thrift protocol.

## Testing
This fix is tested with an integration test that validates that the
correct links are fetched when fetching from a pair of chunk index and
start row offset. There are also unit tests to validate correct client
behaviour when unexpected responses are received from the server.

## Additional Notes to the Reviewer
I also made some changes to the validation of the results. Commented
within the PR.

---------

Co-authored-by: tejassp-db <>
Co-authored-by: Samikshya Chand <148681192+samikshya-db@users.noreply.github.com>

Author: tejassp-db

.github/workflows/prIntegrationTests.yml

@@ -16,7 +16,7 @@ jobs:
         include:
           # SQL_EXEC mode: Tests SEA client behavior
           # Note: CircuitBreakerIntegrationTests requires THRIFT_SERVER mode (tested in second matrix entry)
-          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!M2MAuthIntegrationTests,!CircuitBreakerIntegrationTests
+          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!M2MAuthIntegrationTests,!CircuitBreakerIntegrationTests,!ThriftCloudFetchFakeIntegrationTests
             fake-service-type: 'SQL_EXEC'
           # THRIFT_SERVER mode: Tests Thrift client behavior and circuit breaker fallback
           - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests,!SqlExecApiIntegrationTests

.github/workflows/runIntegrationTests.yml

@@ -2,7 +2,7 @@ name: Integration Tests Workflow - Main Branch
 
 on:
   push:
-    branches: [main]
+    branches: [ main ]
 
 jobs:
   build-and-test:
@@ -14,7 +14,7 @@ jobs:
     strategy:
       matrix:
         include:
-          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!CircuitBreakerIntegrationTests
+          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!CircuitBreakerIntegrationTests,!ThriftCloudFetchFakeIntegrationTests
             token-secret: DATABRICKS_TOKEN
             fake-service-type: 'SQL_EXEC'
           - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests,!SqlExecApiIntegrationTests

NEXT_CHANGELOG.md

@@ -8,6 +8,10 @@
 ### Updated
 
 ### Fixed
+
+- [PECOBLR-1131] Fix incorrect refetching of expired CloudFetch links when using Thrift protocol.
 - Fixed logging to respect params when the driver is shaded.
+
 ---
-*Note: When making changes, please add your change under the appropriate section with a brief description.*
+*Note: When making changes, please add your change under the appropriate section
+with a brief description.*

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java

@@ -6,6 +6,7 @@
 import com.databricks.jdbc.api.IDatabricksResultSet;
 import com.databricks.jdbc.api.IExecutionStatus;
 import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
+import com.databricks.jdbc.api.impl.arrow.ChunkProvider;
 import com.databricks.jdbc.api.impl.converters.ConverterHelper;
 import com.databricks.jdbc.api.impl.converters.ObjectConverter;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
@@ -41,6 +42,7 @@
 import java.util.Calendar;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Supplier;
 import org.apache.http.entity.InputStreamEntity;
 
@@ -2008,6 +2010,14 @@ private BigDecimal applyScaleToBigDecimal(BigDecimal bigDecimal, int columnIndex
     return bigDecimal.setScale(scale, RoundingMode.HALF_UP);
   }
 
+  @VisibleForTesting
+  public Optional<ChunkProvider> getChunkProvider() {
+    if (executionResult instanceof ArrowStreamResult) {
+      return Optional.ofNullable(((ArrowStreamResult) executionResult).getChunkProvider());
+    }
+    return Optional.empty();
+  }
+
   @Override
   public String toString() {
     return (new ToStringer(DatabricksResultSet.class))

src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java

@@ -117,11 +117,11 @@ public Long getChunkIndex() {
   }
 
   /**
-   * Returns the starting row offset for this chunk.
+   * Returns the start row offset of this chunk in the overall result set.
    *
-   * @return the row offset
+   * @return row offset
    */
-  public long getRowOffset() {
+  public long getStartRowOffset() {
     return rowOffset;
   }
 
@@ -156,6 +156,10 @@ public boolean releaseChunk() {
     return true;
   }
 
+  public ExternalLink getChunkLink() {
+    return chunkLink;
+  }
+
   /**
    * Sets the external link details for this chunk.
    *

src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java

@@ -215,6 +215,10 @@ public long getAllowedChunksInMemory() {
     return allowedChunksInMemory;
   }
 
+  public T getChunkByIndex(long chunkIndex) {
+    return chunkIndexToChunksMap.get(chunkIndex);
+  }
+
   /** Subclasses should override this method to perform their specific cleanup. */
   protected void doClose() {
     // Default implementation does nothing

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java

@@ -364,6 +364,16 @@ public long getChunkCount() {
     return chunkProvider.getChunkCount();
   }
 
+  /**
+   * Returns the chunk provider for testing purposes.
+   *
+   * @return the chunk provider
+   */
+  @VisibleForTesting
+  public ChunkProvider getChunkProvider() {
+    return chunkProvider;
+  }
+
   private void setColumnInfo(TGetResultSetMetadataResp resultManifest) {
     columnInfos = new ArrayList<>();
     if (resultManifest.getSchema() == null) {

src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java

@@ -215,15 +215,18 @@ private void triggerNextBatchDownload() {
       return;
     }
 
+    // Calculate row offset for this batch
+    final long batchStartRowOffset = getChunkStartRowOffset(batchStartIndex);
+
     LOGGER.info("Starting batch download from index {}", batchStartIndex);
     currentDownloadTask =
         CompletableFuture.runAsync(
             () -> {
               try {
-                // rowOffset is 0 here as this service is used by RemoteChunkProvider (SEA-only)
-                // which fetches by chunkIndex, not rowOffset
                 ChunkLinkFetchResult result =
-                    session.getDatabricksClient().getResultChunks(statementId, batchStartIndex, 0);
+                    session
+                        .getDatabricksClient()
+                        .getResultChunks(statementId, batchStartIndex, batchStartRowOffset);
                 LOGGER.info(
                     "Retrieved {} links for batch starting at {} for statement id {}",
                     result.getChunkLinks().size(),
@@ -419,6 +422,28 @@ private void prepareNewBatchDownload(long startIndex) {
     isDownloadChainStarted.set(false);
   }
 
+  /**
+   * Gets the start row offset for a given chunk index.
+   *
+   * @param chunkIndex the chunk index to get the row offset for
+   * @return the start row offset for the chunk
+   */
+  private long getChunkStartRowOffset(long chunkIndex) {
+    T chunk = chunkIndexToChunksMap.get(chunkIndex);
+    if (chunk == null) {
+      // Should never happen.
+      throw new IllegalStateException(
+          "Chunk not found in map for index "
+              + chunkIndex
+              + ". "
+              + "Total chunks: "
+              + totalChunks
+              + ", StatementId: "
+              + statementId);
+    }
+    return chunk.getStartRowOffset();
+  }
+
   private boolean isChunkLinkExpired(ExternalLink link) {
     if (link == null || link.getExpiration() == null) {
       LOGGER.warn("Link or expiration is null, assuming link is expired");

src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTask.java

@@ -53,7 +53,7 @@ public Void call() throws DatabricksSQLException {
           if (chunk.isChunkLinkInvalid()) {
             LOGGER.debug("Link invalid for chunk {}, refetching", chunk.getChunkIndex());
             ExternalLink freshLink =
-                linkFetcher.refetchLink(chunk.getChunkIndex(), chunk.getRowOffset());
+                linkFetcher.refetchLink(chunk.getChunkIndex(), chunk.getStartRowOffset());
             chunk.setChunkLink(freshLink);
           }
 

src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java

@@ -94,7 +94,6 @@ public static void verifySuccessStatus(TStatus status, String errorContext, Stri
                   "Error thrift response received [%s] for statementId [%s]",
                   errorContext, statementId)
               : String.format("Error thrift response received [%s]", errorContext);
-      LOGGER.error(errorMessage);
       throw new DatabricksHttpException(errorMessage, status.getSqlState());
     }
   }