Fix Arrow chunk download telemetry error bucketing

NEXT_CHANGELOG.md

@@ -9,6 +9,9 @@
 
 ### Fixed
 - Fixed connections failing when the same parameter is provided in both the JDBC URL and the connection properties, with the JDBC URL taking precedence.
+- Fixed Arrow chunk error handling to avoid exporting internal lifecycle states, preserve parsing
+  failures after download succeeds, and include statement/chunk context on terminal failures.
+
 - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails.
 
 - Throw `DatabricksSQLException` instead of an unchecked `ClassCastException` when a complex-type getter (`getArray`, `getStruct`, `getMap`) is called on a column of a different complex type.

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

@@ -171,9 +171,10 @@ public T getChunk() throws DatabricksSQLException {
           "Operation interrupted while waiting for chunk ready",
           e,
           DatabricksDriverErrorCode.THREAD_INTERRUPTED_ERROR);
-    } catch (ExecutionException | TimeoutException e) {
-      throw new DatabricksSQLException(
-          "Failed to ready chunk", e.getCause(), DatabricksDriverErrorCode.CHUNK_READY_ERROR);
+    } catch (ExecutionException e) {
+      throw createChunkReadyException(e.getCause());
+    } catch (TimeoutException e) {
+      throw createChunkReadyException(e);
     }
     long waitMs = (System.nanoTime() - waitStart) / 1_000_000;
     LOGGER.debug(
@@ -185,6 +186,14 @@ public T getChunk() throws DatabricksSQLException {
     return chunk;
   }
 
+  static DatabricksSQLException createChunkReadyException(Throwable cause) {
+    if (cause instanceof DatabricksSQLException) {
+      return (DatabricksSQLException) cause;
+    }
+    return new DatabricksSQLException(
+        "Failed to ready chunk", cause, DatabricksDriverErrorCode.CHUNK_READY_ERROR);
+  }
+
   /** {@inheritDoc} */
   @Override
   public boolean next() throws DatabricksSQLException {

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

@@ -1,7 +1,7 @@
 package com.databricks.jdbc.api.impl.arrow;
 
 import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createExternalLink;
-import static com.databricks.jdbc.common.util.ValidationUtil.checkHTTPError;
+import static com.databricks.jdbc.common.util.ValidationUtil.checkHTTPErrorWithoutThrowingError;
 import static com.databricks.jdbc.telemetry.TelemetryHelper.getStatementIdString;
 
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
@@ -16,6 +16,7 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowResultLink;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.jdbc.telemetry.TelemetryHelper;
 import com.databricks.sdk.service.sql.BaseChunkInfo;
 import java.io.IOException;
@@ -80,7 +81,10 @@ protected void downloadData(
       addHeaders(getRequest, chunkLink.getHttpHeaders());
       // Retry would be done in http client, we should not bother about that here
       response = httpClient.execute(getRequest, true);
-      checkHTTPError(response);
+      String httpError = checkHTTPErrorWithoutThrowingError(response);
+      if (!httpError.isEmpty()) {
+        throw new IOException(httpError);
+      }
       long downloadTimeMs = (System.nanoTime() - startTime) / 1_000_000;
 
       // Record chunk download latency telemetry
@@ -127,8 +131,10 @@ protected void downloadData(
           readTimeMs - downloadTimeMs,
           decompressTimeMs,
           totalTimeMs);
+    } catch (DatabricksParsingException e) {
+      throw e;
     } catch (Exception e) {
-      handleFailure(e, ChunkStatus.DOWNLOAD_FAILED);
+      handleDownloadFailure(e);
     } finally {
       if (response != null) {
         response.close();
@@ -139,13 +145,13 @@ protected void downloadData(
   /**
    * {@inheritDoc}
    *
-   * <p>Handles failures that occur during chunk download or processing. Sets the error message,
-   * logs the error, updates the chunk status, and throws a DatabricksParsingException.
+   * <p>Handles failures that occur while processing a downloaded chunk. Sets the error message,
+   * logs the error, updates the chunk status, and preserves an existing typed parsing exception or
+   * emits the canonical Arrow parsing error.
    *
    * @param exception the exception that caused the failure
-   * @param failedStatus the status to set for the chunk after failure (e.g. {@link
-   *     ChunkStatus#DOWNLOAD_FAILED} or {@link ChunkStatus#PROCESSING_FAILED})
-   * @throws DatabricksParsingException always thrown with the error message and original exception
+   * @param failedStatus the status to set for the chunk after failure
+   * @throws DatabricksParsingException always thrown; existing typed exceptions are preserved
    */
   @Override
   protected void handleFailure(Exception exception, ChunkStatus failedStatus)
@@ -156,7 +162,24 @@ protected void handleFailure(Exception exception, ChunkStatus failedStatus)
             this.chunkIndex, this.statementId, exception);
     LOGGER.error(this.errorMessage);
     setStatus(failedStatus);
-    throw new DatabricksParsingException(errorMessage, exception, failedStatus.toString());
+    if (exception instanceof DatabricksParsingException) {
+      throw (DatabricksParsingException) exception;
+    }
+    throw new DatabricksParsingException(
+        errorMessage, exception, DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+  }
+
+  private void handleDownloadFailure(Exception exception) throws IOException {
+    errorMessage =
+        String.format(
+            "Data download failed for chunk index [%d] and statement [%s]. Exception [%s]",
+            this.chunkIndex, this.statementId, exception);
+    LOGGER.warn(this.errorMessage);
+    setStatus(ChunkStatus.DOWNLOAD_FAILED);
+    if (exception instanceof IOException) {
+      throw (IOException) exception;
+    }
+    throw new IOException(errorMessage, exception);
   }
 
   private void addHeaders(HttpGet getRequest, Map<String, String> headers) {

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

@@ -3,6 +3,7 @@
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.util.DatabricksThreadContextHolder;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
+import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -40,7 +41,7 @@ class ChunkDownloadTask implements DatabricksCallableTask {
   }
 
   @Override
-  public Void call() throws DatabricksSQLException, ExecutionException, InterruptedException {
+  public Void call() throws DatabricksSQLException {
     int retries = 0;
     boolean downloadSuccessful = false;
 
@@ -80,6 +81,27 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
               chunk.getChunkIndex(),
               taskTotalMs,
               retries);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          throw new DatabricksSQLException(
+              "Interrupted while retrieving chunk download link",
+              e,
+              statementId,
+              chunk.getChunkIndex(),
+              DatabricksDriverErrorCode.THREAD_INTERRUPTED_ERROR.name());
+        } catch (ExecutionException e) {
+          Throwable cause = e.getCause() != null ? e.getCause() : e;
+          if (cause instanceof DatabricksSQLException) {
+            throw (DatabricksSQLException) cause;
+          }
+          throw new DatabricksSQLException(
+              "Failed to retrieve chunk download link",
+              cause,
+              statementId,
+              chunk.getChunkIndex(),
+              DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name());
+        } catch (DatabricksParsingException e) {
+          throw e;
         } catch (IOException | DatabricksSQLException e) {
           retries++;
           if (retries >= MAX_RETRIES) {
@@ -89,7 +111,6 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
                 MAX_RETRIES,
                 chunk.getChunkIndex(),
                 e.getMessage());
-            chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED);
             throw new DatabricksSQLException(
                 "Failed to download chunk after multiple attempts",
                 e,
@@ -125,16 +146,11 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
             "Uncaught exception during chunk download. Chunk index: {}, Error: {}",
             chunk.getChunkIndex(),
             Arrays.toString(uncaughtException.getStackTrace()));
-        // Status is set to DOWNLOAD_SUCCEEDED in the happy path. For any failure case,
-        // explicitly set status to DOWNLOAD_FAILED here to ensure consistent error handling
-        chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED);
-        chunk
-            .getChunkReadyFuture()
-            .completeExceptionally(
-                new DatabricksSQLException(
-                    "Download failed for chunk index " + chunk.getChunkIndex(),
-                    uncaughtException,
-                    DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR));
+        if (chunk.getStatus() != ChunkStatus.DOWNLOAD_FAILED
+            && chunk.getStatus() != ChunkStatus.PROCESSING_FAILED) {
+          chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED);
+        }
+        chunk.getChunkReadyFuture().completeExceptionally(uncaughtException);
       }
 
       DatabricksThreadContextHolder.clearAllContext();

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

@@ -4,6 +4,7 @@
 import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.common.util.DatabricksThreadContextHolder;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
+import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -83,6 +84,8 @@ public Void call() throws DatabricksSQLException {
               taskTotalMs,
               retries);
 
+        } catch (DatabricksParsingException e) {
+          throw e;
         } catch (IOException | SQLException e) {
           retries++;
           if (retries >= MAX_RETRIES) {
@@ -97,7 +100,9 @@ public Void call() throws DatabricksSQLException {
                     "Failed to download chunk %d after %d attempts",
                     chunk.getChunkIndex(), MAX_RETRIES),
                 e,
-                DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR);
+                statementId,
+                chunk.getChunkIndex(),
+                DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name());
           } else {
             LOGGER.warn(
                 "Retry {} for chunk {}: {}", retries, chunk.getChunkIndex(), e.getMessage());
@@ -125,14 +130,11 @@ public Void call() throws DatabricksSQLException {
             "Download failed for chunk {}: {}",
             chunk.getChunkIndex(),
             uncaughtException != null ? uncaughtException.getMessage() : "unknown");
-        chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED);
-        chunk
-            .getChunkReadyFuture()
-            .completeExceptionally(
-                new DatabricksSQLException(
-                    "Download failed for chunk " + chunk.getChunkIndex(),
-                    uncaughtException,
-                    DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR));
+        if (chunk.getStatus() != ChunkStatus.DOWNLOAD_FAILED
+            && chunk.getStatus() != ChunkStatus.PROCESSING_FAILED) {
+          chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED);
+        }
+        chunk.getChunkReadyFuture().completeExceptionally(uncaughtException);
       }
 
       DatabricksThreadContextHolder.clearAllContext();

src/test/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProviderTest.java

@@ -0,0 +1,44 @@
+package com.databricks.jdbc.api.impl.arrow;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import java.util.concurrent.TimeoutException;
+import org.junit.jupiter.api.Test;
+
+public class AbstractRemoteChunkProviderTest {
+
+  @Test
+  void typedChunkFailureIsPreserved() {
+    DatabricksSQLException typedFailure =
+        new DatabricksSQLException(
+            "Arrow parsing failed", DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+
+    DatabricksSQLException result =
+        AbstractRemoteChunkProvider.createChunkReadyException(typedFailure);
+
+    assertSame(typedFailure, result);
+  }
+
+  @Test
+  void untypedChunkFailureIsWrapped() {
+    IllegalStateException cause = new IllegalStateException("unexpected failure");
+
+    DatabricksSQLException result = AbstractRemoteChunkProvider.createChunkReadyException(cause);
+
+    assertEquals(DatabricksDriverErrorCode.CHUNK_READY_ERROR.name(), result.getSQLState());
+    assertSame(cause, result.getCause());
+  }
+
+  @Test
+  void timeoutIsPreservedAsCause() {
+    TimeoutException timeout = new TimeoutException("chunk was not ready");
+
+    DatabricksSQLException result = AbstractRemoteChunkProvider.createChunkReadyException(timeout);
+
+    assertEquals(DatabricksDriverErrorCode.CHUNK_READY_ERROR.name(), result.getSQLState());
+    assertSame(timeout, result.getCause());
+  }
+}