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

.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());
     }
   }

src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java

@@ -128,11 +128,11 @@ DatabricksResultSet getStatementResult(
    * </ul>
    *
    * @param statementId statement-Id for which chunk should be fetched
-   * @param chunkIndex chunkIndex for which chunk should be fetched (used by SEA)
-   * @param rowOffset row offset for fetching results (used by Thrift with FETCH_ABSOLUTE)
-   * @return ChunkLinkFetchResult containing links and continuation information
+   * @param chunkIndex chunkIndex for which chunk should be fetched
+   * @param chunkStartRowOffset the row offset where the chunk starts in the result set
    */
-  ChunkLinkFetchResult getResultChunks(StatementId statementId, long chunkIndex, long rowOffset)
+  ChunkLinkFetchResult getResultChunks(
+      StatementId statementId, long chunkIndex, long chunkStartRowOffset)
       throws DatabricksSQLException;
 
   /**

src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java

@@ -430,8 +430,9 @@ public void cancelStatement(StatementId typedStatementId) throws DatabricksSQLEx
 
   @Override
   public ChunkLinkFetchResult getResultChunks(
-      StatementId typedStatementId, long chunkIndex, long rowOffset) throws DatabricksSQLException {
-    // SEA uses chunkIndex; rowOffset is ignored
+      StatementId typedStatementId, long chunkIndex, long chunkStartRowOffset)
+      throws DatabricksSQLException {
+    DatabricksThreadContextHolder.setStatementId(typedStatementId);
     String statementId = typedStatementId.toSQLExecStatementId();
     LOGGER.debug(
         "getResultChunks(statementId={}, chunkIndex={}) using SEA client", statementId, chunkIndex);