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/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..f1dfb76fe 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,8 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte chunk.getChunkIndex(), taskTotalMs, retries); + } catch (DatabricksParsingException e) { + throw e; } catch (IOException | DatabricksSQLException e) { retries++; if (retries >= MAX_RETRIES) { @@ -89,7 +92,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 +127,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/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..008928a33 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 @@ -56,11 +56,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 +80,42 @@ 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()); + 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()); - 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); + 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); + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> downloadFuture.get()); + 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..29381e2d2 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,32 @@ 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); + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> downloadFuture.get()); + assertSame(thrown, executionException.getCause()); } @Test