diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index 8e44c82b1..6014da699 100644
--- a/NEXT_CHANGELOG.md
+++ b/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.
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java
index 4717852f4..4f2bab626 100644
--- a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java
+++ b/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 {
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java
index fe1d4c7f5..f762e7562 100644
--- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java
+++ b/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}
*
- *
Handles failures that occur during chunk download or processing. Sets the error message,
- * logs the error, updates the chunk status, and throws a DatabricksParsingException.
+ *
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 headers) {
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java
index 217b74a21..7bbf003eb 100644
--- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java
+++ b/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();
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTask.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTask.java
index 177d0bf9f..83f53e7b3 100644
--- a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTask.java
+++ b/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();
diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProviderTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProviderTest.java
new file mode 100644
index 000000000..1227ffd80
--- /dev/null
+++ b/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());
+ }
+}
diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java
index bfc897722..445ecfdbe 100644
--- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java
+++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunkStatusTest.java
@@ -1,7 +1,11 @@
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;
@@ -9,6 +13,7 @@
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,7 +167,7 @@ public java.util.concurrent.Future 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
@@ -123,7 +175,7 @@ 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
diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTaskTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTaskTest.java
index b10336f1e..0c7d3e698 100644
--- a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTaskTest.java
+++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTaskTest.java
@@ -10,6 +10,7 @@
import com.databricks.jdbc.exception.DatabricksSQLException;
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.ByteArrayInputStream;
import java.io.InputStream;
@@ -30,6 +31,7 @@
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
+import org.mockito.MockedStatic;
import org.mockito.MockitoAnnotations;
import org.mockito.junit.jupiter.MockitoExtension;
@@ -56,11 +58,7 @@ void testRetryLogicWithSocketException() throws Exception {
when(chunk.isChunkLinkInvalid()).thenReturn(false);
when(chunk.getChunkIndex()).thenReturn(7L);
when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE);
- DatabricksParsingException throwableError =
- new DatabricksParsingException(
- "Connection reset",
- new SocketException("Connection reset"),
- DatabricksDriverErrorCode.INVALID_STATE);
+ SocketException throwableError = new SocketException("Connection reset");
// Simulate SocketException for the first two attempts, then succeed
doThrow(throwableError)
@@ -84,21 +82,92 @@ void testRetryLogicExhaustedWithSocketException() throws Exception {
when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE);
// Simulate SocketException for all attempts
- doThrow(
- new DatabricksParsingException(
- "Connection reset",
- new SocketException("Connection reset"),
- DatabricksDriverErrorCode.INVALID_STATE))
+ doThrow(new SocketException("Connection reset"))
.when(chunk)
.downloadData(httpClient, CompressionCodec.NONE, 0.1);
- assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call());
- verify(chunk, times(ChunkDownloadTask.MAX_RETRIES))
- .downloadData(httpClient, CompressionCodec.NONE, 0.1);
- assertTrue(downloadFuture.isDone());
+ try (MockedStatic telemetry = mockStatic(TelemetryHelper.class)) {
+ DatabricksSQLException thrown =
+ assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call());
+ assertEquals(DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name(), thrown.getSQLState());
+ verify(chunk, times(ChunkDownloadTask.MAX_RETRIES))
+ .downloadData(httpClient, CompressionCodec.NONE, 0.1);
+ assertTrue(downloadFuture.isDone());
+ ExecutionException executionException =
+ assertThrows(ExecutionException.class, () -> downloadFuture.get());
+ assertSame(thrown, executionException.getCause());
+ assertSame(
+ thrown,
+ AbstractRemoteChunkProvider.createChunkReadyException(executionException.getCause()));
+ telemetry.verify(
+ () ->
+ TelemetryHelper.exportFailureLog(
+ null,
+ DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name(),
+ "Failed to download chunk after multiple attempts",
+ null,
+ 7L,
+ com.databricks.jdbc.common.TelemetryLogLevel.ERROR),
+ times(1));
+ }
+ }
+
+ @Test
+ void testLinkFetchFailureIsReportedAsChunkDownloadError() throws Exception {
+ when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture);
+ when(chunk.isChunkLinkInvalid()).thenReturn(true);
+ when(chunk.getChunkIndex()).thenReturn(7L);
+ CompletableFuture failedLink = new CompletableFuture<>();
+ failedLink.completeExceptionally(new IllegalStateException("link fetch failed"));
+ when(chunkLinkDownloadService.getLinkForChunk(7L)).thenReturn(failedLink);
+
+ try (MockedStatic telemetry = mockStatic(TelemetryHelper.class)) {
+ DatabricksSQLException thrown =
+ assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call());
+
+ assertEquals(DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name(), thrown.getSQLState());
+ ExecutionException executionException =
+ assertThrows(ExecutionException.class, () -> downloadFuture.get());
+ assertSame(thrown, executionException.getCause());
+ assertSame(
+ thrown,
+ AbstractRemoteChunkProvider.createChunkReadyException(executionException.getCause()));
+ verify(chunk, never()).downloadData(any(), any(), anyDouble());
+ telemetry.verify(
+ () ->
+ TelemetryHelper.exportFailureLog(
+ null,
+ DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name(),
+ "Failed to retrieve chunk download link",
+ null,
+ 7L,
+ com.databricks.jdbc.common.TelemetryLogLevel.ERROR),
+ times(1));
+ }
+ }
+
+ @Test
+ void testProcessingFailureIsNotRetried() throws Exception {
+ when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture);
+ when(chunk.isChunkLinkInvalid()).thenReturn(false);
+ when(chunk.getChunkIndex()).thenReturn(7L);
+ when(chunk.getStatus()).thenReturn(ChunkStatus.PROCESSING_FAILED);
+ when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE);
+ DatabricksParsingException processingError =
+ new DatabricksParsingException(
+ "Arrow parsing failed", DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+ doThrow(processingError).when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1);
+
+ DatabricksParsingException thrown =
+ assertThrows(DatabricksParsingException.class, () -> chunkDownloadTask.call());
+
+ assertSame(processingError, thrown);
+ verify(chunk, times(1)).downloadData(httpClient, CompressionCodec.NONE, 0.1);
+ verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY);
+ verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_FAILED);
ExecutionException executionException =
assertThrows(ExecutionException.class, () -> downloadFuture.get());
- assertInstanceOf(DatabricksSQLException.class, executionException.getCause());
+ assertSame(thrown, executionException.getCause());
}
@Test
diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTaskTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTaskTest.java
index 4d2e5aed1..4cbce970a 100644
--- a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTaskTest.java
+++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTaskTest.java
@@ -77,11 +77,7 @@ void testRetryLogicWithSocketException() throws Exception {
when(chunk.isChunkLinkInvalid()).thenReturn(false);
when(chunk.getChunkIndex()).thenReturn(7L);
- DatabricksParsingException throwableError =
- new DatabricksParsingException(
- "Connection reset",
- new SocketException("Connection reset"),
- DatabricksDriverErrorCode.INVALID_STATE);
+ SocketException throwableError = new SocketException("Connection reset");
// Simulate SocketException for the first two attempts, then succeed
doThrow(throwableError)
@@ -106,15 +102,13 @@ void testRetryLogicExhaustedWithSocketException() throws Exception {
when(chunk.getChunkIndex()).thenReturn(7L);
// Simulate SocketException for all attempts
- doThrow(
- new DatabricksParsingException(
- "Connection reset",
- new SocketException("Connection reset"),
- DatabricksDriverErrorCode.INVALID_STATE))
+ doThrow(new SocketException("Connection reset"))
.when(chunk)
.downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
- assertThrows(DatabricksSQLException.class, () -> downloadTask.call());
+ DatabricksSQLException thrown =
+ assertThrows(DatabricksSQLException.class, () -> downloadTask.call());
+ assertEquals(DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name(), thrown.getSQLState());
// Should attempt MAX_RETRIES (5) times
verify(chunk, times(5))
@@ -124,7 +118,33 @@ void testRetryLogicExhaustedWithSocketException() throws Exception {
ExecutionException executionException =
assertThrows(ExecutionException.class, () -> downloadFuture.get());
- assertInstanceOf(DatabricksSQLException.class, executionException.getCause());
+ assertSame(thrown, executionException.getCause());
+ }
+
+ @Test
+ void testProcessingFailureIsNotRetried() throws Exception {
+ when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture);
+ when(chunk.isChunkLinkInvalid()).thenReturn(false);
+ when(chunk.getChunkIndex()).thenReturn(7L);
+ when(chunk.getStatus()).thenReturn(ChunkStatus.PROCESSING_FAILED);
+ DatabricksParsingException processingError =
+ new DatabricksParsingException(
+ "Arrow parsing failed", DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+ doThrow(processingError)
+ .when(chunk)
+ .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
+
+ DatabricksParsingException thrown =
+ assertThrows(DatabricksParsingException.class, () -> downloadTask.call());
+
+ assertSame(processingError, thrown);
+ verify(chunk, times(1))
+ .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
+ verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY);
+ verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_FAILED);
+ ExecutionException executionException =
+ assertThrows(ExecutionException.class, () -> downloadFuture.get());
+ assertSame(thrown, executionException.getCause());
}
@Test