From add47ed9991725feb4ea7ff9ca9bb3f4e16325cf Mon Sep 17 00:00:00 2001 From: Prathamesh Baviskar Date: Thu, 20 Aug 2026 09:25:09 +0000 Subject: [PATCH 1/2] [PECOBLR-3984] Implement fix Signed-off-by: Prathamesh Baviskar --- .../impl/arrow/AbstractArrowResultChunk.java | 19 ++++- .../jdbc/api/impl/arrow/ArrowResultChunk.java | 29 +++++++- .../api/impl/arrow/ChunkDownloadTask.java | 68 ++++++++++++++++-- .../arrow/StreamingChunkDownloadTask.java | 68 ++++++++++++++++-- .../jdbc/common/util/ValidationUtil.java | 3 +- .../impl/http/DatabricksHttpRetryHandler.java | 2 +- .../exception/DatabricksHttpException.java | 24 +++++++ .../api/impl/arrow/ChunkDownloadTaskTest.java | 63 ++++++++++++++++ .../arrow/StreamingChunkDownloadTaskTest.java | 71 +++++++++++++++++++ 9 files changed, 328 insertions(+), 19 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java index 17fa6f2531..0a353ff237 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java @@ -6,6 +6,7 @@ import com.databricks.jdbc.common.util.DriverUtil; 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.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; @@ -187,13 +188,27 @@ public ChunkStatus getStatus() { * @param compressionCodec the compression codec to use for decompression * @param speedThreshold the minimum expected download speed in MB/s for logging warnings * @throws DatabricksParsingException if there is an error parsing the data + * @throws DatabricksHttpException if the server returns an HTTP error response * @throws IOException if there is an error downloading or reading the data */ protected abstract void downloadData( IDatabricksHttpClient httpClient, CompressionCodec compressionCodec, double speedThreshold) - throws DatabricksParsingException, IOException; + throws DatabricksParsingException, DatabricksHttpException, IOException; - /** Handles a failure during the download or processing of this chunk. */ + /** + * Handles a failure during the download or processing of this chunk. + * + *

Implementations must set the chunk's {@link #errorMessage}, log the error, update the chunk + * status, and throw a {@link DatabricksParsingException}. When the root cause is a {@link + * com.databricks.jdbc.exception.DatabricksHttpException}, the HTTP status code should be included + * in the error message so that callers (retry loops) can log it without unwrapping the cause + * chain. + * + * @param exception the exception that caused the failure + * @param failedStatus the status to set for the chunk ({@link ChunkStatus#DOWNLOAD_FAILED} or + * {@link ChunkStatus#PROCESSING_FAILED}) + * @throws DatabricksParsingException always thrown; wraps the original exception + */ protected abstract void handleFailure(Exception exception, ChunkStatus failedStatus) throws DatabricksParsingException; 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 fe1d4c7f58..6daf19fe02 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 @@ -10,6 +10,7 @@ import com.databricks.jdbc.common.util.DecompressionUtil; 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.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; @@ -71,7 +72,7 @@ public static Builder builder() { @Override protected void downloadData( IDatabricksHttpClient httpClient, CompressionCodec compressionCodec, double speedThreshold) - throws DatabricksParsingException, IOException { + throws DatabricksParsingException, DatabricksHttpException, IOException { CloseableHttpResponse response = null; long startTime = System.nanoTime(); try { @@ -147,13 +148,35 @@ protected void downloadData( * ChunkStatus#DOWNLOAD_FAILED} or {@link ChunkStatus#PROCESSING_FAILED}) * @throws DatabricksParsingException always thrown with the error message and original exception */ + /** + * {@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. + * + *

When the exception is a {@link DatabricksHttpException} carrying an HTTP status code, the + * status code is appended to the error message (e.g. {@code [HTTP 403]}) so it is visible in + * retry logs without requiring callers to unwrap the cause chain. + * + * @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 + */ @Override protected void handleFailure(Exception exception, ChunkStatus failedStatus) throws DatabricksParsingException { + String httpStatusAnnotation = ""; + if (exception instanceof DatabricksHttpException) { + int httpStatus = ((DatabricksHttpException) exception).getHttpStatusCode(); + if (httpStatus > 0) { + httpStatusAnnotation = String.format(" [HTTP %d]", httpStatus); + } + } errorMessage = String.format( - "Data parsing failed for chunk index [%d] and statement [%s]. Exception [%s]", - this.chunkIndex, this.statementId, exception); + "Data parsing failed for chunk index [%d] and statement [%s]%s. Exception [%s]", + this.chunkIndex, this.statementId, httpStatusAnnotation, exception); LOGGER.error(this.errorMessage); setStatus(failedStatus); throw new DatabricksParsingException(errorMessage, exception, failedStatus.toString()); 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 217b74a212..2e0da1ecc7 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,8 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; +import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler; +import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -17,7 +19,6 @@ class ChunkDownloadTask implements DatabricksCallableTask { private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ChunkDownloadTask.class); public static final int MAX_RETRIES = 5; - private static final long RETRY_DELAY_MS = 1500; // 1.5 seconds private final ArrowResultChunk chunk; private final IDatabricksHttpClient httpClient; private final ChunkDownloadManager chunkDownloader; @@ -81,29 +82,50 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte taskTotalMs, retries); } catch (IOException | DatabricksSQLException e) { + int httpStatus = extractHttpStatus(e); retries++; if (retries >= MAX_RETRIES) { LOGGER.error( e, - "Failed to download chunk after %d attempts. Chunk index: %d, Error: %s", + "Failed to download chunk after %d attempts. Chunk index: %d, HTTP status: %d, Error: %s", MAX_RETRIES, chunk.getChunkIndex(), + httpStatus, e.getMessage()); chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); throw new DatabricksSQLException( - "Failed to download chunk after multiple attempts", + String.format( + "Failed to download chunk after multiple attempts (HTTP status: %d)", + httpStatus), + e, + statementId, + chunk.getChunkIndex(), + DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name()); + } else if (isPermanentHttpFailure(httpStatus)) { + LOGGER.error( + e, + "Permanent HTTP %d error for chunk index: %d, will not retry. Error: %s", + httpStatus, + chunk.getChunkIndex(), + e.getMessage()); + chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); + throw new DatabricksSQLException( + String.format( + "Permanent HTTP %d error downloading chunk %d", + httpStatus, chunk.getChunkIndex()), e, statementId, chunk.getChunkIndex(), DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name()); } else { + long delayMs = DatabricksHttpRetryHandler.calculateExponentialBackoff(retries); LOGGER.warn( String.format( - "Retry attempt %d for chunk index: %d, Error: %s", - retries, chunk.getChunkIndex(), e.getMessage())); + "Retry attempt %d for chunk index: %d, HTTP status: %d, retryDelayMs: %d, Error: %s", + retries, chunk.getChunkIndex(), httpStatus, delayMs, e.getMessage())); chunk.setStatus(ChunkStatus.DOWNLOAD_RETRY); try { - Thread.sleep(RETRY_DELAY_MS); + Thread.sleep(delayMs); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new DatabricksSQLException( @@ -142,4 +164,38 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte return null; } + + /** + * Extracts the HTTP status code from the exception or its direct cause when either is a {@link + * DatabricksHttpException}. Returns 0 when the failure is a network error with no HTTP response. + */ + private static int extractHttpStatus(Exception e) { + if (e instanceof DatabricksHttpException) { + return ((DatabricksHttpException) e).getHttpStatusCode(); + } + Throwable cause = e.getCause(); + if (cause instanceof DatabricksHttpException) { + return ((DatabricksHttpException) cause).getHttpStatusCode(); + } + return 0; + } + + /** + * Returns {@code true} for 4xx HTTP status codes that represent permanent client errors not worth + * retrying. Excludes codes that are transient or recoverable: + * + *

+ */ + private static boolean isPermanentHttpFailure(int httpStatus) { + return httpStatus >= 400 + && httpStatus < 500 + && httpStatus != 403 + && httpStatus != 408 + && httpStatus != 429; + } } 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 177d0bf9f0..d0d5c8c60d 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,8 @@ import com.databricks.jdbc.common.CompressionCodec; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; +import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler; +import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -23,7 +25,6 @@ public class StreamingChunkDownloadTask implements Callable { JdbcLoggerFactory.getLogger(StreamingChunkDownloadTask.class); private static final int MAX_RETRIES = 5; - private static final long RETRY_DELAY_MS = 1500; private final ArrowResultChunk chunk; private final IDatabricksHttpClient httpClient; @@ -84,26 +85,47 @@ public Void call() throws DatabricksSQLException { retries); } catch (IOException | SQLException e) { + int httpStatus = extractHttpStatus(e); retries++; if (retries >= MAX_RETRIES) { LOGGER.error( - "Failed to download chunk {} after {} attempts: {}", + "Failed to download chunk {} after {} attempts: HTTP status: {}, Error: {}", chunk.getChunkIndex(), MAX_RETRIES, + httpStatus, e.getMessage()); // Status set to DOWNLOAD_FAILED in the finally block throw new DatabricksSQLException( String.format( - "Failed to download chunk %d after %d attempts", - chunk.getChunkIndex(), MAX_RETRIES), + "Failed to download chunk %d after %d attempts (HTTP status: %d)", + chunk.getChunkIndex(), MAX_RETRIES, httpStatus), + e, + DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR); + } else if (isPermanentHttpFailure(httpStatus)) { + LOGGER.error( + "Permanent HTTP {} error for chunk {}, will not retry. Error: {}", + httpStatus, + chunk.getChunkIndex(), + e.getMessage()); + chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); + throw new DatabricksSQLException( + String.format( + "Permanent HTTP %d error downloading chunk %d", + httpStatus, chunk.getChunkIndex()), e, DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR); } else { + long delayMs = DatabricksHttpRetryHandler.calculateExponentialBackoff(retries); LOGGER.warn( - "Retry {} for chunk {}: {}", retries, chunk.getChunkIndex(), e.getMessage()); + "Retry {} for chunk {}: HTTP status: {}, retryDelayMs: {}, Error: {}", + retries, + chunk.getChunkIndex(), + httpStatus, + delayMs, + e.getMessage()); chunk.setStatus(ChunkStatus.DOWNLOAD_RETRY); try { - Thread.sleep(RETRY_DELAY_MS); + Thread.sleep(delayMs); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new DatabricksSQLException( @@ -140,4 +162,38 @@ public Void call() throws DatabricksSQLException { return null; } + + /** + * Extracts the HTTP status code from the exception or its direct cause when either is a {@link + * DatabricksHttpException}. Returns 0 when the failure is a network error with no HTTP response. + */ + private static int extractHttpStatus(Exception e) { + if (e instanceof DatabricksHttpException) { + return ((DatabricksHttpException) e).getHttpStatusCode(); + } + Throwable cause = e.getCause(); + if (cause instanceof DatabricksHttpException) { + return ((DatabricksHttpException) cause).getHttpStatusCode(); + } + return 0; + } + + /** + * Returns {@code true} for 4xx HTTP status codes that represent permanent client errors not worth + * retrying. Excludes codes that are transient or recoverable: + * + *
    + *
  • 403: pre-signed URL may have expired; {@code isChunkLinkInvalid()} will detect and + * refresh the link on the next iteration. + *
  • 408: request timeout — transient. + *
  • 429: rate limit — transient. + *
+ */ + private static boolean isPermanentHttpFailure(int httpStatus) { + return httpStatus >= 400 + && httpStatus < 500 + && httpStatus != 403 + && httpStatus != 408 + && httpStatus != 429; + } } diff --git a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java index 4c29953100..31c27bdbba 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java @@ -133,12 +133,13 @@ public static String checkHTTPErrorWithoutThrowingError(HttpResponse response) { public static void checkHTTPError(HttpResponse response) throws DatabricksHttpException, IOException { + int statusCode = response.getStatusLine().getStatusCode(); String errorReason = checkHTTPErrorWithoutThrowingError(response); if (errorReason.equals(EMPTY_STRING)) { return; } LOGGER.error(errorReason); - throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + throw new DatabricksHttpException(errorReason, statusCode, DEFAULT_HTTP_EXCEPTION_SQLSTATE); } /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java index 94650ef19c..bccfe71e47 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java @@ -242,7 +242,7 @@ static long calculateDelayInMillis(int errorCode, int executionCount, int retryI } } - static long calculateExponentialBackoff(int executionCount) { + public static long calculateExponentialBackoff(int executionCount) { return Math.min( MIN_BACKOFF_INTERVAL * (long) Math.pow(DEFAULT_BACKOFF_FACTOR, executionCount), MAX_RETRY_INTERVAL); diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java index 3bb04c8d6e..163befc2bc 100644 --- a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java @@ -5,20 +5,44 @@ /** Exception class to handle http errors while downloading chunk data from external links. */ public class DatabricksHttpException extends DatabricksSQLException { + /** HTTP response status code; 0 when no HTTP response was received (e.g. network error). */ + private final int httpStatusCode; + public DatabricksHttpException( String message, Throwable cause, DatabricksDriverErrorCode sqlCode) { super(message, cause, sqlCode); + this.httpStatusCode = 0; } public DatabricksHttpException(String message, DatabricksDriverErrorCode internalCode) { super(message, null, internalCode.toString()); + this.httpStatusCode = 0; } public DatabricksHttpException(String message, String sqlState) { super(message, null, sqlState); + this.httpStatusCode = 0; + } + + /** + * Creates an HTTP exception carrying the response status code for programmatic differentiation of + * transient vs. permanent failures. + */ + public DatabricksHttpException(String message, int httpStatusCode, String sqlState) { + super(message, null, sqlState); + this.httpStatusCode = httpStatusCode; } public DatabricksHttpException(String message, Throwable throwable, String sqlState) { super(message, throwable, sqlState); + this.httpStatusCode = 0; + } + + /** + * Returns the HTTP response status code associated with this exception, or 0 if no HTTP response + * was received (e.g. connection reset before a response arrived). + */ + public int getHttpStatusCode() { + return httpStatusCode; } } 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 b10336f1e3..0ca1172814 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 @@ -6,6 +6,7 @@ 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.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ExternalLink; @@ -208,6 +209,68 @@ void testRetryLogicWithRealChunkAndStatusTransitions() throws Exception { verify(spiedChunk, times(1)).initializeData(any(InputStream.class)); } + @Test + void testRetryOnTransient500HttpError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(3L); + when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); + + DatabricksHttpException http500 = + new DatabricksHttpException("HTTP request failed by code: 500", 500, "08000"); + + // Fail with a 500 on the first attempt, then succeed + doThrow(http500).doNothing().when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); + + chunkDownloadTask.call(); + + verify(chunk, times(2)).downloadData(httpClient, CompressionCodec.NONE, 0.1); + verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); + assertTrue(downloadFuture.isDone()); + assertDoesNotThrow(() -> downloadFuture.get()); + } + + @Test + void testFailFastOnPermanent404HttpError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(4L); + when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); + + DatabricksHttpException http404 = + new DatabricksHttpException("HTTP request failed by code: 404", 404, "08000"); + + doThrow(http404).when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); + + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call()); + assertTrue(thrown.getMessage().contains("404"), "Error message should contain HTTP status 404"); + // Single download attempt — permanent failures are not retried + verify(chunk, times(1)).downloadData(httpClient, CompressionCodec.NONE, 0.1); + verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY); + } + + @Test + void testRetryOn403LinkExpiredError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(5L); + when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); + + DatabricksHttpException http403 = + new DatabricksHttpException("HTTP request failed by code: 403", 403, "08000"); + + // 403 is retryable (pre-signed URL may have expired); succeed on second attempt + doThrow(http403).doNothing().when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); + + chunkDownloadTask.call(); + + verify(chunk, times(2)).downloadData(httpClient, CompressionCodec.NONE, 0.1); + verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); + assertTrue(downloadFuture.isDone()); + assertDoesNotThrow(() -> downloadFuture.get()); + } + private BaseChunkInfo createMockBaseChunkInfo(long chunkIndex, long rowCount, long rowOffset) { BaseChunkInfo mockChunkInfo = mock(BaseChunkInfo.class); when(mockChunkInfo.getChunkIndex()).thenReturn(chunkIndex); 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 4d2e5aed10..3aac69b588 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 @@ -6,6 +6,7 @@ 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.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ExternalLink; @@ -180,6 +181,76 @@ void testLinkRefreshFailureAndRetry() throws Exception { assertDoesNotThrow(() -> downloadFuture.get()); } + @Test + void testRetryOnTransient500HttpError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(3L); + + DatabricksHttpException http500 = + new DatabricksHttpException("HTTP request failed by code: 500", 500, "08000"); + + // Fail with a 500 on the first attempt, then succeed + doThrow(http500) + .doNothing() + .when(chunk) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + + downloadTask.call(); + + verify(chunk, times(2)) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); + assertTrue(downloadFuture.isDone()); + assertDoesNotThrow(() -> downloadFuture.get()); + } + + @Test + void testFailFastOnPermanent404HttpError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(4L); + + DatabricksHttpException http404 = + new DatabricksHttpException("HTTP request failed by code: 404", 404, "08000"); + + doThrow(http404) + .when(chunk) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + + DatabricksSQLException thrown = + assertThrows(DatabricksSQLException.class, () -> downloadTask.call()); + assertTrue(thrown.getMessage().contains("404"), "Error message should contain HTTP status 404"); + // Single download attempt — permanent failures are not retried + verify(chunk, times(1)) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY); + } + + @Test + void testRetryOn403LinkExpiredError() throws Exception { + when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); + when(chunk.isChunkLinkInvalid()).thenReturn(false); + when(chunk.getChunkIndex()).thenReturn(5L); + + DatabricksHttpException http403 = + new DatabricksHttpException("HTTP request failed by code: 403", 403, "08000"); + + // 403 is retryable (pre-signed URL may have expired); succeed on second attempt + doThrow(http403) + .doNothing() + .when(chunk) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + + downloadTask.call(); + + verify(chunk, times(2)) + .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); + verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); + assertTrue(downloadFuture.isDone()); + assertDoesNotThrow(() -> downloadFuture.get()); + } + @Test void testStatusTransitionsDuringRetries() throws Exception { StatementId statementId = new StatementId("test-statement-123"); From 70dc71075b99585e54b5ed58f00788c30334899c Mon Sep 17 00:00:00 2001 From: Prathamesh Baviskar Date: Thu, 20 Aug 2026 14:38:40 +0000 Subject: [PATCH 2/2] [PECOBLR-3984] Add jitter to chunk download retries Signed-off-by: Prathamesh Baviskar --- NEXT_CHANGELOG.md | 2 + .../impl/arrow/AbstractArrowResultChunk.java | 19 +--- .../jdbc/api/impl/arrow/ArrowResultChunk.java | 29 +----- .../api/impl/arrow/ChunkDownloadTask.java | 79 ++++------------ .../jdbc/api/impl/arrow/ChunkRetryPolicy.java | 29 ++++++ .../arrow/StreamingChunkDownloadTask.java | 84 ++++++----------- .../jdbc/common/util/ValidationUtil.java | 3 +- .../impl/http/DatabricksHttpRetryHandler.java | 2 +- .../exception/DatabricksHttpException.java | 24 ----- .../api/impl/arrow/ChunkDownloadTaskTest.java | 77 +++------------- .../api/impl/arrow/ChunkRetryPolicyTest.java | 28 ++++++ .../arrow/StreamingChunkDownloadTaskTest.java | 90 ++++--------------- 12 files changed, 139 insertions(+), 327 deletions(-) create mode 100644 src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicy.java create mode 100644 src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicyTest.java diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index e12c8875f4..e956b1b17d 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,8 @@ - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. ### Fixed +- Added jitter to chunk-download retries to reduce synchronized retry bursts during transient failures without increasing the existing maximum delay. + - 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/AbstractArrowResultChunk.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java index 0a353ff237..17fa6f2531 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java @@ -6,7 +6,6 @@ import com.databricks.jdbc.common.util.DriverUtil; 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.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; @@ -188,27 +187,13 @@ public ChunkStatus getStatus() { * @param compressionCodec the compression codec to use for decompression * @param speedThreshold the minimum expected download speed in MB/s for logging warnings * @throws DatabricksParsingException if there is an error parsing the data - * @throws DatabricksHttpException if the server returns an HTTP error response * @throws IOException if there is an error downloading or reading the data */ protected abstract void downloadData( IDatabricksHttpClient httpClient, CompressionCodec compressionCodec, double speedThreshold) - throws DatabricksParsingException, DatabricksHttpException, IOException; + throws DatabricksParsingException, IOException; - /** - * Handles a failure during the download or processing of this chunk. - * - *

Implementations must set the chunk's {@link #errorMessage}, log the error, update the chunk - * status, and throw a {@link DatabricksParsingException}. When the root cause is a {@link - * com.databricks.jdbc.exception.DatabricksHttpException}, the HTTP status code should be included - * in the error message so that callers (retry loops) can log it without unwrapping the cause - * chain. - * - * @param exception the exception that caused the failure - * @param failedStatus the status to set for the chunk ({@link ChunkStatus#DOWNLOAD_FAILED} or - * {@link ChunkStatus#PROCESSING_FAILED}) - * @throws DatabricksParsingException always thrown; wraps the original exception - */ + /** Handles a failure during the download or processing of this chunk. */ protected abstract void handleFailure(Exception exception, ChunkStatus failedStatus) throws DatabricksParsingException; 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 6daf19fe02..fe1d4c7f58 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 @@ -10,7 +10,6 @@ import com.databricks.jdbc.common.util.DecompressionUtil; 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.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; @@ -72,7 +71,7 @@ public static Builder builder() { @Override protected void downloadData( IDatabricksHttpClient httpClient, CompressionCodec compressionCodec, double speedThreshold) - throws DatabricksParsingException, DatabricksHttpException, IOException { + throws DatabricksParsingException, IOException { CloseableHttpResponse response = null; long startTime = System.nanoTime(); try { @@ -148,35 +147,13 @@ protected void downloadData( * ChunkStatus#DOWNLOAD_FAILED} or {@link ChunkStatus#PROCESSING_FAILED}) * @throws DatabricksParsingException always thrown with the error message and original exception */ - /** - * {@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. - * - *

When the exception is a {@link DatabricksHttpException} carrying an HTTP status code, the - * status code is appended to the error message (e.g. {@code [HTTP 403]}) so it is visible in - * retry logs without requiring callers to unwrap the cause chain. - * - * @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 - */ @Override protected void handleFailure(Exception exception, ChunkStatus failedStatus) throws DatabricksParsingException { - String httpStatusAnnotation = ""; - if (exception instanceof DatabricksHttpException) { - int httpStatus = ((DatabricksHttpException) exception).getHttpStatusCode(); - if (httpStatus > 0) { - httpStatusAnnotation = String.format(" [HTTP %d]", httpStatus); - } - } errorMessage = String.format( - "Data parsing failed for chunk index [%d] and statement [%s]%s. Exception [%s]", - this.chunkIndex, this.statementId, httpStatusAnnotation, exception); + "Data parsing failed for chunk index [%d] and statement [%s]. Exception [%s]", + this.chunkIndex, this.statementId, exception); LOGGER.error(this.errorMessage); setStatus(failedStatus); throw new DatabricksParsingException(errorMessage, exception, failedStatus.toString()); 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 2e0da1ecc7..9fca6dcf87 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,8 +3,6 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; -import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler; -import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -25,6 +23,7 @@ class ChunkDownloadTask implements DatabricksCallableTask { private final IDatabricksConnectionContext connectionContext; private final String statementId; private final ChunkLinkDownloadService linkDownloadService; + private final ChunkRetryPolicy retryPolicy; Throwable uncaughtException = null; ChunkDownloadTask( @@ -32,12 +31,22 @@ class ChunkDownloadTask implements DatabricksCallableTask { IDatabricksHttpClient httpClient, ChunkDownloadManager chunkDownloader, ChunkLinkDownloadService linkDownloadService) { + this(chunk, httpClient, chunkDownloader, linkDownloadService, new ChunkRetryPolicy()); + } + + ChunkDownloadTask( + ArrowResultChunk chunk, + IDatabricksHttpClient httpClient, + ChunkDownloadManager chunkDownloader, + ChunkLinkDownloadService linkDownloadService, + ChunkRetryPolicy retryPolicy) { this.chunk = chunk; this.httpClient = httpClient; this.chunkDownloader = chunkDownloader; this.connectionContext = DatabricksThreadContextHolder.getConnectionContext(); this.statementId = DatabricksThreadContextHolder.getStatementId(); this.linkDownloadService = linkDownloadService; + this.retryPolicy = retryPolicy; } @Override @@ -82,50 +91,30 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte taskTotalMs, retries); } catch (IOException | DatabricksSQLException e) { - int httpStatus = extractHttpStatus(e); retries++; if (retries >= MAX_RETRIES) { LOGGER.error( e, - "Failed to download chunk after %d attempts. Chunk index: %d, HTTP status: %d, Error: %s", + "Failed to download chunk after %d attempts. Chunk index: %d, Error: %s", MAX_RETRIES, chunk.getChunkIndex(), - httpStatus, - e.getMessage()); - chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); - throw new DatabricksSQLException( - String.format( - "Failed to download chunk after multiple attempts (HTTP status: %d)", - httpStatus), - e, - statementId, - chunk.getChunkIndex(), - DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name()); - } else if (isPermanentHttpFailure(httpStatus)) { - LOGGER.error( - e, - "Permanent HTTP %d error for chunk index: %d, will not retry. Error: %s", - httpStatus, - chunk.getChunkIndex(), e.getMessage()); chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); throw new DatabricksSQLException( - String.format( - "Permanent HTTP %d error downloading chunk %d", - httpStatus, chunk.getChunkIndex()), + "Failed to download chunk after multiple attempts", e, statementId, chunk.getChunkIndex(), DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name()); } else { - long delayMs = DatabricksHttpRetryHandler.calculateExponentialBackoff(retries); + long retryDelayMs = retryPolicy.getRetryDelayMs(); LOGGER.warn( String.format( - "Retry attempt %d for chunk index: %d, HTTP status: %d, retryDelayMs: %d, Error: %s", - retries, chunk.getChunkIndex(), httpStatus, delayMs, e.getMessage())); + "Retry attempt %d for chunk index: %d, retryDelayMs: %d, Error: %s", + retries, chunk.getChunkIndex(), retryDelayMs, e.getMessage())); chunk.setStatus(ChunkStatus.DOWNLOAD_RETRY); try { - Thread.sleep(delayMs); + retryPolicy.sleep(retryDelayMs); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new DatabricksSQLException( @@ -164,38 +153,4 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte return null; } - - /** - * Extracts the HTTP status code from the exception or its direct cause when either is a {@link - * DatabricksHttpException}. Returns 0 when the failure is a network error with no HTTP response. - */ - private static int extractHttpStatus(Exception e) { - if (e instanceof DatabricksHttpException) { - return ((DatabricksHttpException) e).getHttpStatusCode(); - } - Throwable cause = e.getCause(); - if (cause instanceof DatabricksHttpException) { - return ((DatabricksHttpException) cause).getHttpStatusCode(); - } - return 0; - } - - /** - * Returns {@code true} for 4xx HTTP status codes that represent permanent client errors not worth - * retrying. Excludes codes that are transient or recoverable: - * - *

    - *
  • 403: pre-signed URL may have expired; {@code isChunkLinkInvalid()} will detect and - * refresh the link on the next iteration. - *
  • 408: request timeout — transient. - *
  • 429: rate limit — transient. - *
- */ - private static boolean isPermanentHttpFailure(int httpStatus) { - return httpStatus >= 400 - && httpStatus < 500 - && httpStatus != 403 - && httpStatus != 408 - && httpStatus != 429; - } } diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicy.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicy.java new file mode 100644 index 0000000000..ddecf73ef6 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicy.java @@ -0,0 +1,29 @@ +package com.databricks.jdbc.api.impl.arrow; + +import java.util.concurrent.ThreadLocalRandom; +import java.util.function.LongUnaryOperator; + +/** Adds equal jitter so concurrent chunk downloads do not retry in lockstep. */ +final class ChunkRetryPolicy { + + static final long MIN_RETRY_DELAY_MS = 750; + static final long MAX_RETRY_DELAY_MS = 1500; + private final LongUnaryOperator randomLong; + + ChunkRetryPolicy() { + this(bound -> ThreadLocalRandom.current().nextLong(bound)); + } + + ChunkRetryPolicy(LongUnaryOperator randomLong) { + this.randomLong = randomLong; + } + + long getRetryDelayMs() { + long range = MAX_RETRY_DELAY_MS - MIN_RETRY_DELAY_MS + 1; + return MIN_RETRY_DELAY_MS + randomLong.applyAsLong(range); + } + + void sleep(long retryDelayMs) throws InterruptedException { + Thread.sleep(retryDelayMs); + } +} 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 d0d5c8c60d..c318fb4102 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,8 +4,6 @@ import com.databricks.jdbc.common.CompressionCodec; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; -import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler; -import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; @@ -31,6 +29,7 @@ public class StreamingChunkDownloadTask implements Callable { private final CompressionCodec compressionCodec; private final LinkRefresher linkRefresher; private final double cloudFetchSpeedThreshold; + private final ChunkRetryPolicy retryPolicy; // Capture caller's thread context for telemetry/logging on the download thread private final IDatabricksConnectionContext connectionContext; @@ -42,11 +41,28 @@ public StreamingChunkDownloadTask( CompressionCodec compressionCodec, LinkRefresher linkRefresher, double cloudFetchSpeedThreshold) { + this( + chunk, + httpClient, + compressionCodec, + linkRefresher, + cloudFetchSpeedThreshold, + new ChunkRetryPolicy()); + } + + StreamingChunkDownloadTask( + ArrowResultChunk chunk, + IDatabricksHttpClient httpClient, + CompressionCodec compressionCodec, + LinkRefresher linkRefresher, + double cloudFetchSpeedThreshold, + ChunkRetryPolicy retryPolicy) { this.chunk = chunk; this.httpClient = httpClient; this.compressionCodec = compressionCodec; this.linkRefresher = linkRefresher; this.cloudFetchSpeedThreshold = cloudFetchSpeedThreshold; + this.retryPolicy = retryPolicy; this.connectionContext = DatabricksThreadContextHolder.getConnectionContext(); this.statementId = DatabricksThreadContextHolder.getStatementId(); } @@ -85,47 +101,31 @@ public Void call() throws DatabricksSQLException { retries); } catch (IOException | SQLException e) { - int httpStatus = extractHttpStatus(e); retries++; if (retries >= MAX_RETRIES) { LOGGER.error( - "Failed to download chunk {} after {} attempts: HTTP status: {}, Error: {}", + "Failed to download chunk {} after {} attempts: {}", chunk.getChunkIndex(), MAX_RETRIES, - httpStatus, e.getMessage()); // Status set to DOWNLOAD_FAILED in the finally block throw new DatabricksSQLException( String.format( - "Failed to download chunk %d after %d attempts (HTTP status: %d)", - chunk.getChunkIndex(), MAX_RETRIES, httpStatus), - e, - DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR); - } else if (isPermanentHttpFailure(httpStatus)) { - LOGGER.error( - "Permanent HTTP {} error for chunk {}, will not retry. Error: {}", - httpStatus, - chunk.getChunkIndex(), - e.getMessage()); - chunk.setStatus(ChunkStatus.DOWNLOAD_FAILED); - throw new DatabricksSQLException( - String.format( - "Permanent HTTP %d error downloading chunk %d", - httpStatus, chunk.getChunkIndex()), + "Failed to download chunk %d after %d attempts", + chunk.getChunkIndex(), MAX_RETRIES), e, DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR); } else { - long delayMs = DatabricksHttpRetryHandler.calculateExponentialBackoff(retries); + long retryDelayMs = retryPolicy.getRetryDelayMs(); LOGGER.warn( - "Retry {} for chunk {}: HTTP status: {}, retryDelayMs: {}, Error: {}", + "Retry {} for chunk {} in {} ms: {}", retries, chunk.getChunkIndex(), - httpStatus, - delayMs, + retryDelayMs, e.getMessage()); chunk.setStatus(ChunkStatus.DOWNLOAD_RETRY); try { - Thread.sleep(delayMs); + retryPolicy.sleep(retryDelayMs); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw new DatabricksSQLException( @@ -162,38 +162,4 @@ public Void call() throws DatabricksSQLException { return null; } - - /** - * Extracts the HTTP status code from the exception or its direct cause when either is a {@link - * DatabricksHttpException}. Returns 0 when the failure is a network error with no HTTP response. - */ - private static int extractHttpStatus(Exception e) { - if (e instanceof DatabricksHttpException) { - return ((DatabricksHttpException) e).getHttpStatusCode(); - } - Throwable cause = e.getCause(); - if (cause instanceof DatabricksHttpException) { - return ((DatabricksHttpException) cause).getHttpStatusCode(); - } - return 0; - } - - /** - * Returns {@code true} for 4xx HTTP status codes that represent permanent client errors not worth - * retrying. Excludes codes that are transient or recoverable: - * - *
    - *
  • 403: pre-signed URL may have expired; {@code isChunkLinkInvalid()} will detect and - * refresh the link on the next iteration. - *
  • 408: request timeout — transient. - *
  • 429: rate limit — transient. - *
- */ - private static boolean isPermanentHttpFailure(int httpStatus) { - return httpStatus >= 400 - && httpStatus < 500 - && httpStatus != 403 - && httpStatus != 408 - && httpStatus != 429; - } } diff --git a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java index 31c27bdbba..4c29953100 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java @@ -133,13 +133,12 @@ public static String checkHTTPErrorWithoutThrowingError(HttpResponse response) { public static void checkHTTPError(HttpResponse response) throws DatabricksHttpException, IOException { - int statusCode = response.getStatusLine().getStatusCode(); String errorReason = checkHTTPErrorWithoutThrowingError(response); if (errorReason.equals(EMPTY_STRING)) { return; } LOGGER.error(errorReason); - throw new DatabricksHttpException(errorReason, statusCode, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE); } /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java index bccfe71e47..94650ef19c 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpRetryHandler.java @@ -242,7 +242,7 @@ static long calculateDelayInMillis(int errorCode, int executionCount, int retryI } } - public static long calculateExponentialBackoff(int executionCount) { + static long calculateExponentialBackoff(int executionCount) { return Math.min( MIN_BACKOFF_INTERVAL * (long) Math.pow(DEFAULT_BACKOFF_FACTOR, executionCount), MAX_RETRY_INTERVAL); diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java index 163befc2bc..3bb04c8d6e 100644 --- a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java @@ -5,44 +5,20 @@ /** Exception class to handle http errors while downloading chunk data from external links. */ public class DatabricksHttpException extends DatabricksSQLException { - /** HTTP response status code; 0 when no HTTP response was received (e.g. network error). */ - private final int httpStatusCode; - public DatabricksHttpException( String message, Throwable cause, DatabricksDriverErrorCode sqlCode) { super(message, cause, sqlCode); - this.httpStatusCode = 0; } public DatabricksHttpException(String message, DatabricksDriverErrorCode internalCode) { super(message, null, internalCode.toString()); - this.httpStatusCode = 0; } public DatabricksHttpException(String message, String sqlState) { super(message, null, sqlState); - this.httpStatusCode = 0; - } - - /** - * Creates an HTTP exception carrying the response status code for programmatic differentiation of - * transient vs. permanent failures. - */ - public DatabricksHttpException(String message, int httpStatusCode, String sqlState) { - super(message, null, sqlState); - this.httpStatusCode = httpStatusCode; } public DatabricksHttpException(String message, Throwable throwable, String sqlState) { super(message, throwable, sqlState); - this.httpStatusCode = 0; - } - - /** - * Returns the HTTP response status code associated with this exception, or 0 if no HTTP response - * was received (e.g. connection reset before a response arrived). - */ - public int getHttpStatusCode() { - return httpStatusCode; } } 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 0ca1172814..5d64c45b30 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 @@ -6,7 +6,6 @@ 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.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ExternalLink; @@ -40,6 +39,7 @@ public class ChunkDownloadTaskTest { @Mock IDatabricksHttpClient httpClient; @Mock RemoteChunkProvider remoteChunkProvider; @Mock ChunkLinkDownloadService chunkLinkDownloadService; + @Mock ChunkRetryPolicy retryPolicy; private ChunkDownloadTask chunkDownloadTask; private CompletableFuture downloadFuture; @@ -48,7 +48,8 @@ void setUp() { MockitoAnnotations.openMocks(this); downloadFuture = new CompletableFuture<>(); chunkDownloadTask = - new ChunkDownloadTask(chunk, httpClient, remoteChunkProvider, chunkLinkDownloadService); + new ChunkDownloadTask( + chunk, httpClient, remoteChunkProvider, chunkLinkDownloadService, retryPolicy); } @Test @@ -57,6 +58,7 @@ void testRetryLogicWithSocketException() throws Exception { when(chunk.isChunkLinkInvalid()).thenReturn(false); when(chunk.getChunkIndex()).thenReturn(7L); when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); + when(retryPolicy.getRetryDelayMs()).thenReturn(750L, 1500L); DatabricksParsingException throwableError = new DatabricksParsingException( "Connection reset", @@ -73,6 +75,9 @@ void testRetryLogicWithSocketException() throws Exception { chunkDownloadTask.call(); verify(chunk, times(3)).downloadData(httpClient, CompressionCodec.NONE, 0.1); + verify(retryPolicy, times(2)).getRetryDelayMs(); + verify(retryPolicy).sleep(750); + verify(retryPolicy).sleep(1500); assertTrue(downloadFuture.isDone()); assertDoesNotThrow(() -> downloadFuture.get()); } @@ -83,6 +88,7 @@ void testRetryLogicExhaustedWithSocketException() throws Exception { when(chunk.isChunkLinkInvalid()).thenReturn(false); when(chunk.getChunkIndex()).thenReturn(7L); when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); + when(retryPolicy.getRetryDelayMs()).thenReturn(750L, 1000L, 1250L, 1500L); // Simulate SocketException for all attempts doThrow( @@ -96,6 +102,11 @@ void testRetryLogicExhaustedWithSocketException() throws Exception { assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call()); verify(chunk, times(ChunkDownloadTask.MAX_RETRIES)) .downloadData(httpClient, CompressionCodec.NONE, 0.1); + verify(retryPolicy, times(ChunkDownloadTask.MAX_RETRIES - 1)).getRetryDelayMs(); + verify(retryPolicy).sleep(750); + verify(retryPolicy).sleep(1000); + verify(retryPolicy).sleep(1250); + verify(retryPolicy).sleep(1500); assertTrue(downloadFuture.isDone()); ExecutionException executionException = assertThrows(ExecutionException.class, () -> downloadFuture.get()); @@ -209,68 +220,6 @@ void testRetryLogicWithRealChunkAndStatusTransitions() throws Exception { verify(spiedChunk, times(1)).initializeData(any(InputStream.class)); } - @Test - void testRetryOnTransient500HttpError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(3L); - when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); - - DatabricksHttpException http500 = - new DatabricksHttpException("HTTP request failed by code: 500", 500, "08000"); - - // Fail with a 500 on the first attempt, then succeed - doThrow(http500).doNothing().when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); - - chunkDownloadTask.call(); - - verify(chunk, times(2)).downloadData(httpClient, CompressionCodec.NONE, 0.1); - verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); - assertTrue(downloadFuture.isDone()); - assertDoesNotThrow(() -> downloadFuture.get()); - } - - @Test - void testFailFastOnPermanent404HttpError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(4L); - when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); - - DatabricksHttpException http404 = - new DatabricksHttpException("HTTP request failed by code: 404", 404, "08000"); - - doThrow(http404).when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); - - DatabricksSQLException thrown = - assertThrows(DatabricksSQLException.class, () -> chunkDownloadTask.call()); - assertTrue(thrown.getMessage().contains("404"), "Error message should contain HTTP status 404"); - // Single download attempt — permanent failures are not retried - verify(chunk, times(1)).downloadData(httpClient, CompressionCodec.NONE, 0.1); - verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY); - } - - @Test - void testRetryOn403LinkExpiredError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(5L); - when(remoteChunkProvider.getCompressionCodec()).thenReturn(CompressionCodec.NONE); - - DatabricksHttpException http403 = - new DatabricksHttpException("HTTP request failed by code: 403", 403, "08000"); - - // 403 is retryable (pre-signed URL may have expired); succeed on second attempt - doThrow(http403).doNothing().when(chunk).downloadData(httpClient, CompressionCodec.NONE, 0.1); - - chunkDownloadTask.call(); - - verify(chunk, times(2)).downloadData(httpClient, CompressionCodec.NONE, 0.1); - verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); - assertTrue(downloadFuture.isDone()); - assertDoesNotThrow(() -> downloadFuture.get()); - } - private BaseChunkInfo createMockBaseChunkInfo(long chunkIndex, long rowCount, long rowOffset) { BaseChunkInfo mockChunkInfo = mock(BaseChunkInfo.class); when(mockChunkInfo.getChunkIndex()).thenReturn(chunkIndex); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicyTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicyTest.java new file mode 100644 index 0000000000..3165c72895 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkRetryPolicyTest.java @@ -0,0 +1,28 @@ +package com.databricks.jdbc.api.impl.arrow; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.Test; + +class ChunkRetryPolicyTest { + + @Test + void retryDelayUsesInclusiveEqualJitterBounds() { + AtomicLong requestedBound = new AtomicLong(); + ChunkRetryPolicy minimumPolicy = + new ChunkRetryPolicy( + bound -> { + requestedBound.set(bound); + return 0; + }); + ChunkRetryPolicy maximumPolicy = new ChunkRetryPolicy(bound -> bound - 1); + + long minimumDelay = minimumPolicy.getRetryDelayMs(); + long maximumDelay = maximumPolicy.getRetryDelayMs(); + + assertEquals(751, requestedBound.get()); + assertEquals(ChunkRetryPolicy.MIN_RETRY_DELAY_MS, minimumDelay); + assertEquals(ChunkRetryPolicy.MAX_RETRY_DELAY_MS, maximumDelay); + } +} 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 3aac69b588..824f0d7c0b 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 @@ -6,7 +6,6 @@ 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.exception.DatabricksSQLException; import com.databricks.jdbc.model.core.ExternalLink; @@ -40,6 +39,7 @@ public class StreamingChunkDownloadTaskTest { @Mock private ArrowResultChunk chunk; @Mock private IDatabricksHttpClient httpClient; @Mock private LinkRefresher linkRefresher; + @Mock private ChunkRetryPolicy retryPolicy; private StreamingChunkDownloadTask downloadTask; private CompletableFuture downloadFuture; @@ -49,7 +49,12 @@ void setUp() { downloadFuture = new CompletableFuture<>(); downloadTask = new StreamingChunkDownloadTask( - chunk, httpClient, CompressionCodec.NONE, linkRefresher, CLOUD_FETCH_SPEED_THRESHOLD); + chunk, + httpClient, + CompressionCodec.NONE, + linkRefresher, + CLOUD_FETCH_SPEED_THRESHOLD, + retryPolicy); } @Test @@ -68,6 +73,7 @@ void testSuccessfulDownloadOnFirstAttempt() throws Exception { verify(chunk, times(1)) .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY); + verifyNoInteractions(retryPolicy); assertTrue(downloadFuture.isDone()); assertDoesNotThrow(() -> downloadFuture.get()); } @@ -77,6 +83,7 @@ void testRetryLogicWithSocketException() throws Exception { when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); when(chunk.isChunkLinkInvalid()).thenReturn(false); when(chunk.getChunkIndex()).thenReturn(7L); + when(retryPolicy.getRetryDelayMs()).thenReturn(750L, 1500L); DatabricksParsingException throwableError = new DatabricksParsingException( @@ -96,6 +103,9 @@ void testRetryLogicWithSocketException() throws Exception { verify(chunk, times(3)) .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); verify(chunk, times(2)).setStatus(ChunkStatus.DOWNLOAD_RETRY); + verify(retryPolicy, times(2)).getRetryDelayMs(); + verify(retryPolicy).sleep(750); + verify(retryPolicy).sleep(1500); assertTrue(downloadFuture.isDone()); assertDoesNotThrow(() -> downloadFuture.get()); } @@ -105,6 +115,7 @@ void testRetryLogicExhaustedWithSocketException() throws Exception { when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); when(chunk.isChunkLinkInvalid()).thenReturn(false); when(chunk.getChunkIndex()).thenReturn(7L); + when(retryPolicy.getRetryDelayMs()).thenReturn(750L, 1000L, 1250L, 1500L); // Simulate SocketException for all attempts doThrow( @@ -121,6 +132,11 @@ void testRetryLogicExhaustedWithSocketException() throws Exception { verify(chunk, times(5)) .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_FAILED); + verify(retryPolicy, times(4)).getRetryDelayMs(); + verify(retryPolicy).sleep(750); + verify(retryPolicy).sleep(1000); + verify(retryPolicy).sleep(1250); + verify(retryPolicy).sleep(1500); assertTrue(downloadFuture.isDone()); ExecutionException executionException = @@ -181,76 +197,6 @@ void testLinkRefreshFailureAndRetry() throws Exception { assertDoesNotThrow(() -> downloadFuture.get()); } - @Test - void testRetryOnTransient500HttpError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(3L); - - DatabricksHttpException http500 = - new DatabricksHttpException("HTTP request failed by code: 500", 500, "08000"); - - // Fail with a 500 on the first attempt, then succeed - doThrow(http500) - .doNothing() - .when(chunk) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - - downloadTask.call(); - - verify(chunk, times(2)) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); - assertTrue(downloadFuture.isDone()); - assertDoesNotThrow(() -> downloadFuture.get()); - } - - @Test - void testFailFastOnPermanent404HttpError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(4L); - - DatabricksHttpException http404 = - new DatabricksHttpException("HTTP request failed by code: 404", 404, "08000"); - - doThrow(http404) - .when(chunk) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - - DatabricksSQLException thrown = - assertThrows(DatabricksSQLException.class, () -> downloadTask.call()); - assertTrue(thrown.getMessage().contains("404"), "Error message should contain HTTP status 404"); - // Single download attempt — permanent failures are not retried - verify(chunk, times(1)) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - verify(chunk, never()).setStatus(ChunkStatus.DOWNLOAD_RETRY); - } - - @Test - void testRetryOn403LinkExpiredError() throws Exception { - when(chunk.getChunkReadyFuture()).thenReturn(downloadFuture); - when(chunk.isChunkLinkInvalid()).thenReturn(false); - when(chunk.getChunkIndex()).thenReturn(5L); - - DatabricksHttpException http403 = - new DatabricksHttpException("HTTP request failed by code: 403", 403, "08000"); - - // 403 is retryable (pre-signed URL may have expired); succeed on second attempt - doThrow(http403) - .doNothing() - .when(chunk) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - - downloadTask.call(); - - verify(chunk, times(2)) - .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD); - verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY); - assertTrue(downloadFuture.isDone()); - assertDoesNotThrow(() -> downloadFuture.get()); - } - @Test void testStatusTransitionsDuringRetries() throws Exception { StatementId statementId = new StatementId("test-statement-123");