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