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 download telemetry to emit one canonical `CHUNK_DOWNLOAD_ERROR` after retries
+  are exhausted instead of exporting internal lifecycle states for individual attempts.
+
 - 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;
@@ -80,6 +81,19 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
               chunk.getChunkIndex(),
               taskTotalMs,
               retries);
+        } 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 +103,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 +138,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) {
@@ -125,14 +128,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());
+  }
+}

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

@@ -1,14 +1,19 @@
 package com.databricks.jdbc.api.impl.arrow;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
 
 import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.jdbc.telemetry.latency.TelemetryCollectorManager;
 import java.io.ByteArrayInputStream;
 import java.io.FilterInputStream;
@@ -62,14 +67,56 @@ void readError_setsDownloadFailed_notDownloadSucceeded() {
     InputStream erroring = new ErrorInputStream(new ByteArrayInputStream(payload));
     IDatabricksHttpClient http = httpWithEntity(erroring, payload.length);
 
-    // Act + Assert: downloadData should throw parsing exception and status should be
-    // DOWNLOAD_FAILED
-    assertThrows(
-        DatabricksParsingException.class,
-        () -> chunk.downloadData(http, CompressionCodec.NONE, 0.0));
+    assertThrows(IOException.class, () -> chunk.downloadData(http, CompressionCodec.NONE, 0.0));
     assertEquals(ChunkStatus.DOWNLOAD_FAILED, chunk.getStatus());
   }
 
+  @Test
+  void httpError_isReportedAsRetryableDownloadFailure() {
+    byte[] payload = "service unavailable".getBytes();
+    ArrowResultChunk chunk = newChunk();
+    IDatabricksHttpClient http =
+        httpWithEntity(new ByteArrayInputStream(payload), payload.length, 503);
+
+    assertThrows(IOException.class, () -> chunk.downloadData(http, CompressionCodec.NONE, 0.0));
+    assertEquals(ChunkStatus.DOWNLOAD_FAILED, chunk.getStatus());
+  }
+
+  @Test
+  void processingError_isNotReportedAsDownloadError() {
+    byte[] payload = "not an Arrow stream".getBytes();
+    ArrowResultChunk chunk = newChunk();
+    IDatabricksHttpClient http = httpWithEntity(new ByteArrayInputStream(payload), payload.length);
+
+    DatabricksParsingException exception =
+        assertThrows(
+            DatabricksParsingException.class,
+            () -> chunk.downloadData(http, CompressionCodec.NONE, 0.0));
+
+    assertEquals(ChunkStatus.PROCESSING_FAILED, chunk.getStatus());
+    assertEquals(
+        DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(), exception.getSQLState());
+  }
+
+  @Test
+  void typedProcessingError_isPreserved() throws Exception {
+    byte[] payload = "downloaded data".getBytes();
+    ArrowResultChunk chunk = spy(newChunk());
+    IDatabricksHttpClient http = httpWithEntity(new ByteArrayInputStream(payload), payload.length);
+    DatabricksParsingException processingError =
+        new DatabricksParsingException(
+            "typed processing error", DatabricksDriverErrorCode.DECOMPRESSION_ERROR);
+    doThrow(processingError).when(chunk).initializeData(any(InputStream.class));
+
+    DatabricksParsingException thrown =
+        assertThrows(
+            DatabricksParsingException.class,
+            () -> chunk.downloadData(http, CompressionCodec.NONE, 0.0));
+
+    assertSame(processingError, thrown);
+    assertEquals(ChunkStatus.PROCESSING_FAILED, chunk.getStatus());
+  }
+
   private static ArrowResultChunk newChunk() {
     StatementId statementId = new StatementId("stmt-status-test");
     ArrowResultChunk chunk;
@@ -91,18 +138,23 @@ private static ArrowResultChunk newChunk() {
   }
 
   private static IDatabricksHttpClient httpWithEntity(InputStream content, long length) {
+    return httpWithEntity(content, length, 200);
+  }
+
+  private static IDatabricksHttpClient httpWithEntity(
+      InputStream content, long length, int statusCode) {
     return new IDatabricksHttpClient() {
       @Override
       public CloseableHttpResponse execute(org.apache.http.client.methods.HttpUriRequest request)
           throws DatabricksHttpException {
-        return response(content, length);
+        return response(content, length, statusCode);
       }
 
       @Override
       public CloseableHttpResponse execute(
           org.apache.http.client.methods.HttpUriRequest request, boolean supportGzipEncoding)
           throws DatabricksHttpException {
-        return response(content, length);
+        return response(content, length, statusCode);
       }
 
       @Override
@@ -115,15 +167,15 @@ public <T> java.util.concurrent.Future<T> executeAsync(
     };
   }
 
-  private static CloseableHttpResponse response(InputStream content, long length) {
+  private static CloseableHttpResponse response(InputStream content, long length, int statusCode) {
     HttpEntity entity = new InputStreamEntity(content, length);
     return new CloseableHttpResponse() {
       @Override
       public void close() {}
 
       @Override
       public StatusLine getStatusLine() {
-        return new BasicStatusLine(new ProtocolVersion("HTTP", 1, 1), 200, "OK");
+        return new BasicStatusLine(new ProtocolVersion("HTTP", 1, 1), statusCode, "status");
       }
 
       @Override