diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index e12c8875f..297c62e8e 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,7 @@ - `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 +- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/502/504); permanent HTTP errors (401/403/404, and 500) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline, and `CloseOperation`/`CancelOperation` use a tighter retry budget so cleanup on shutdown cannot stall. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. - 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/common/util/ValidationUtil.java b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java index 4c2995310..24547fef9 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java @@ -138,7 +138,8 @@ public static void checkHTTPError(HttpResponse response) return; } LOGGER.error(errorReason); - throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + int statusCode = response.getStatusLine().getStatusCode(); + throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE, statusCode); } /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java index edc85e454..0c2c1242a 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java @@ -72,6 +72,25 @@ public void checkTimeout() throws DatabricksTimeoutException { } } + /** + * Returns the time remaining, in milliseconds, before {@link #checkTimeout()} would actually + * enforce the deadline. Returns {@link Long#MAX_VALUE} when no timeout is configured ({@code + * timeoutSeconds <= 0}), and may return a negative value once the deadline has passed. + * + *

{@code checkTimeout()} compares whole (truncated) seconds and throws only once elapsed time + * exceeds {@code timeoutSeconds}, i.e. at {@code (timeoutSeconds + 1) * 1000} ms. This + * method reports the remaining time to that same enforcement point (not the nominal + * {@code timeoutSeconds * 1000}), so a caller capping a backoff sleep on it does not collapse the + * sleep to zero — and then busy-spin — during the sub-second window before enforcement fires. + */ + public long getRemainingMillis() { + if (timeoutSeconds <= 0) { + return Long.MAX_VALUE; + } + long elapsedMillis = System.currentTimeMillis() - startTimeMillis; + return TimeUnit.SECONDS.toMillis(timeoutSeconds + 1L) - elapsedMillis; + } + /** * Factory method to create a timeout handler for a databricks client with a statement ID. This * works with any client that implements {@link IDatabricksClient}. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java index 966750225..de0c8e1a9 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java @@ -198,7 +198,10 @@ private static void throwHttpException(Exception e, HttpUriRequest request) "Caught error while executing http request: [%s]. Error Message: [%s]", RequestSanitizer.sanitizeRequest(request), e); LOGGER.error(e, errorMsg); - throw new DatabricksHttpException(errorMsg, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + // Preserve the original cause (typically a connection-level IOException) so callers can + // distinguish a genuine transport failure from an HTTP-status error and decide whether the + // operation is safe to retry. + throw new DatabricksHttpException(errorMsg, e, DEFAULT_HTTP_EXCEPTION_SQLSTATE); } @VisibleForTesting diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 18c126b2b..4474fbbbb 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -28,9 +28,13 @@ import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; +import com.google.common.annotations.VisibleForTesting; +import java.io.IOException; import java.sql.SQLException; import java.util.Arrays; import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import org.apache.http.HttpException; import org.apache.thrift.TBase; @@ -53,6 +57,27 @@ final class DatabricksThriftAccessor { TExecuteStatementResp._Fields.OPERATION_HANDLE.getThriftFieldId(); private static final short statusFieldId = TExecuteStatementResp._Fields.STATUS.getThriftFieldId(); + // Bounded, jittered retry for transient transport-level failures (stale pooled connection, + // connection reset, load-balancer idle drop) on idempotent status / close / cancel RPCs. A + // status poll is read-only, and closing or cancelling an operation is idempotent, so repeating + // any of them produces no additional server-side side effects. This is why statement submission + // is deliberately NOT routed through the retry path: re-sending an ExecuteStatement could run the + // query twice. + private static final int MAX_POLL_TRANSPORT_RETRIES = 5; + // Cleanup RPCs (CloseOperation / CancelOperation) run on close / cancel / timeout paths that must + // not hang: one transient blip is worth reconnecting past, but a sustained outage should surface + // fast rather than stall shutdown. Hence a far smaller budget than the status-poll path. + private static final int MAX_CLEANUP_TRANSPORT_RETRIES = 2; + private static final long TRANSPORT_RETRY_MIN_BACKOFF_MILLIS = 1_000L; + private static final long TRANSPORT_RETRY_MAX_BACKOFF_MILLIS = 16_000L; + // Transient HTTP gateway codes that the shared DatabricksHttpRetryHandler does NOT itself retry + // (it only retries 429/503 + configured custom codes). A poll that hits one of these received a + // real HTTP response but from a transiently-unhealthy hop, so re-polling on a fresh connection is + // safe and worthwhile. 429/503 are excluded (owned by the HTTP layer — re-retrying would multiply + // load on a recovering endpoint); 500 is excluded too, as it is typically a genuine server-side + // error rather than a transient hop failure and should surface rather than burn the retry budget. + private static final Set RETRYABLE_TRANSPORT_HTTP_CODES = Set.of(408, 502, 504); + private DatabricksConfig databricksConfig; private final boolean enableDirectResults; private final int asyncPollIntervalMillis; @@ -163,7 +188,10 @@ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, long startR TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksHttpException { try { - return getThriftClient().CancelOperation(req); + return withCleanupTransportRetry( + "CancelOperation", + loggableOperationHandle(req.getOperationHandle()), + () -> getThriftClient().CancelOperation(req)); } catch (TException e) { String errorMessage = String.format( @@ -176,7 +204,10 @@ TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksH TCloseOperationResp closeOperation(TCloseOperationReq req) throws DatabricksHttpException { try { - return getThriftClient().CloseOperation(req); + return withCleanupTransportRetry( + "CloseOperation", + loggableOperationHandle(req.getOperationHandle()), + () -> getThriftClient().CloseOperation(req)); } catch (TException e) { String errorMessage = String.format( @@ -341,9 +372,10 @@ private TGetOperationStatusResp pollTillOperationFinished( timeoutHandler.checkTimeout(); // TTransportException means a transport-level failure (e.g. HTTP 502 Bad Gateway) - // after retries were exhausted. Other TException subtypes propagate unchanged. + // after retries were exhausted. Other TException subtypes propagate unchanged. The timeout + // handler is threaded in so the retry backoff cannot overshoot the statement's queryTimeout. try { - statusResp = getOperationStatus(statusReq, statementId); + statusResp = getOperationStatus(statusReq, statementId, timeoutHandler); } catch (TTransportException e) { throw buildTransportFailureException(statementId.toSQLExecStatementId(), e); } @@ -778,7 +810,12 @@ TFetchResultsResp fetchMetadataResults(TResp response, String contextDescription while (shouldContinuePolling(statusResp)) { metadataTimeoutHandler.checkTimeout(); try { - statusResp = getThriftClient().GetOperationStatus(statusReq); + statusResp = + withTransportRetry( + "GetOperationStatus", + statementId, + metadataTimeoutHandler, + () -> getThriftClient().GetOperationStatus(statusReq)); } catch (TTransportException e) { throw buildTransportFailureException(statementId, e); } @@ -962,6 +999,222 @@ private DatabricksSQLException buildTransportFailureException( return new DatabricksSQLException(errorMsg, e, COMMUNICATION_LINK_FAILURE_SQLSTATE); } + /** A Thrift RPC that is safe to repeat after a transport-level failure. */ + @FunctionalInterface + private interface TransportSafeRpc { + T call() throws TException; + } + + /** Null-safe operation-handle rendering for log lines (handles may be absent). */ + private static String loggableOperationHandle(TOperationHandle operationHandle) { + return operationHandle != null + ? StatementId.loggableStatementId(operationHandle) + : "unknown"; + } + + /** + * Executes an idempotent status-poll RPC with no query-timeout budget (used by the Thrift + * heartbeat / metadata-less status checks). Retries are bounded by {@link + * #MAX_POLL_TRANSPORT_RETRIES}. See {@link #withTransportRetry(String, String, TimeoutHandler, + * int, TransportSafeRpc)} for the full contract. + */ + private T withTransportRetry(String rpcName, String statementId, TransportSafeRpc rpc) + throws TException { + return withoutTimeout(rpcName, statementId, MAX_POLL_TRANSPORT_RETRIES, rpc); + } + + /** + * Executes a cleanup RPC ({@code CloseOperation} / {@code CancelOperation}) with a deliberately + * small retry budget ({@link #MAX_CLEANUP_TRANSPORT_RETRIES}). Cleanup runs on close / cancel / + * timeout paths that must not hang: a single transient blip is worth reconnecting past, but during + * a sustained outage extra retries only delay shutdown, so the budget is far tighter than the + * status-poll path. + */ + private T withCleanupTransportRetry( + String rpcName, String statementId, TransportSafeRpc rpc) throws TException { + return withoutTimeout(rpcName, statementId, MAX_CLEANUP_TRANSPORT_RETRIES, rpc); + } + + /** Shared no-timeout entry point; adapts the deadline-aware core for callers with no deadline. */ + private T withoutTimeout( + String rpcName, String statementId, int maxRetries, TransportSafeRpc rpc) + throws TException { + try { + return withTransportRetry( + rpcName, statementId, /* timeoutHandler= */ null, maxRetries, rpc); + } catch (DatabricksTimeoutException e) { + // Unreachable: a null timeout handler never enforces a deadline. Guard defensively so the + // checked timeout type cannot silently widen this method's contract. + throw new IllegalStateException("Unexpected timeout without an active timeout handler", e); + } + } + + /** + * Deadline-aware status-poll retry with the default poll budget ({@link + * #MAX_POLL_TRANSPORT_RETRIES}). See {@link #withTransportRetry(String, String, TimeoutHandler, + * int, TransportSafeRpc)}. + */ + private T withTransportRetry( + String rpcName, String statementId, TimeoutHandler timeoutHandler, TransportSafeRpc rpc) + throws TException, DatabricksTimeoutException { + return withTransportRetry( + rpcName, statementId, timeoutHandler, MAX_POLL_TRANSPORT_RETRIES, rpc); + } + + /** + * Executes an idempotent Thrift RPC, transparently retrying transient transport-level + * failures on a fresh connection with jittered exponential backoff. + * + *

Every invocation of {@code rpc} builds a new transport, so a retry naturally leases a + * different pooled connection while the broken one is discarded. This lets a still-running + * server-side operation be re-polled — or a completed one be re-closed / re-cancelled — instead + * of being abandoned after a single stale-connection blip and left to expire on the server's + * inactivity timeout. Only RPCs that are safe to repeat may use this path (status polling, + * operation close, cancel); statement submission must not. + * + *

Only failures classified as transient by {@link #isRetryableTransportFailure} are retried: + * genuine connection-level errors (stale pooled connection, reset, socket timeout) and transient + * HTTP gateway codes ({@link #RETRYABLE_TRANSPORT_HTTP_CODES}). Permanent HTTP errors (401/403/404 + * …) and anything the shared {@link + * com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler} already retried and + * exhausted (429/503/custom, which surface with a {@link DatabricksRetryHandlerException} in their + * cause chain) are rethrown on the first attempt — the latter avoids stacking a second retry + * storm on top of the HTTP layer's. + * + *

When {@code timeoutHandler} is non-null the operation's deadline is enforced before each + * backoff sleep and the sleep is capped to the remaining budget, so a failing RPC cannot overshoot + * the statement's {@code queryTimeout}. Retries are bounded by {@code maxRetries}; once exhausted + * the original {@link TTransportException} is rethrown so existing caller-side handling still + * applies. A thread interrupt during a backoff sleep restores the interrupt flag and aborts the + * retry loop. + */ + private T withTransportRetry( + String rpcName, + String statementId, + TimeoutHandler timeoutHandler, + int maxRetries, + TransportSafeRpc rpc) + throws TException, DatabricksTimeoutException { + int attempt = 0; + long backoffMillis = TRANSPORT_RETRY_MIN_BACKOFF_MILLIS; + while (true) { + try { + return rpc.call(); + } catch (TTransportException e) { + if (!isRetryableTransportFailure(e)) { + // Permanent error (e.g. 401/403/404) or one the HTTP layer already retried (429/503): + // surface immediately instead of hanging through the backoff schedule. + throw e; + } + if (++attempt > maxRetries) { + LOGGER.error( + "Transport failure on {} for statement [{}] still failing after {} retries; giving" + + " up. Cause: {}", + rpcName, + statementId, + maxRetries, + e.getMessage()); + throw e; + } + // Enforce the query deadline before sleeping so a failing RPC cannot overshoot the + // statement's queryTimeout by the backoff schedule (may run the timeout action and throw). + if (timeoutHandler != null) { + timeoutHandler.checkTimeout(); + } + // Full-jitter backoff around the current exponential ceiling, spreading concurrent + // reconnect attempts so they do not thunder against a recovering endpoint. + long sleepMillis = + ThreadLocalRandom.current().nextLong(backoffMillis / 2 + 1, backoffMillis + 1); + if (timeoutHandler != null) { + // Cap to the time left before the deadline is actually enforced so we neither overshoot + // nor collapse to a zero-length (busy-spin) sleep in the sub-second window before it. + long remainingMillis = timeoutHandler.getRemainingMillis(); + if (remainingMillis < sleepMillis) { + sleepMillis = Math.max(0L, remainingMillis); + } + } + LOGGER.warn( + "Transport failure on {} for statement [{}] (attempt {}/{}); reconnecting and retrying" + + " in {} ms. Cause: {}", + rpcName, + statementId, + attempt, + maxRetries, + sleepMillis, + e.getMessage()); + try { + backoffSleep(sleepMillis); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw e; + } + backoffMillis = Math.min(backoffMillis * 2, TRANSPORT_RETRY_MAX_BACKOFF_MILLIS); + } + } + } + + /** + * Classifies a transport failure as transient (safe to retry on a fresh connection) or not. + * + *

Because the Thrift transport routes through {@link + * com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClient}, every failure arrives wrapped as a + * {@link TTransportException} whose cause is normally a {@link DatabricksHttpException}. The + * decision: + * + *

+ */ + private static boolean isRetryableTransportFailure(TTransportException e) { + Throwable cause = e.getCause(); + if (cause == null) { + return true; + } + if (chainContains(cause, DatabricksRetryHandlerException.class)) { + return false; + } + if (cause instanceof DatabricksHttpException) { + int statusCode = ((DatabricksHttpException) cause).getStatusCode(); + if (statusCode != 0) { + return RETRYABLE_TRANSPORT_HTTP_CODES.contains(statusCode); + } + // No HTTP response was received: retry only when a real connection-level IOException is the + // underlying cause (guards against unrelated status-less DatabricksHttpExceptions). + return chainContains(cause.getCause(), IOException.class); + } + return cause instanceof IOException; + } + + /** Sleeps for the transport-retry backoff. Extracted as a test seam to keep retry tests fast. */ + @VisibleForTesting + void backoffSleep(long millis) throws InterruptedException { + TimeUnit.MILLISECONDS.sleep(millis); + } + + /** Returns true if {@code throwable} or any exception in its cause chain is of {@code type}. */ + private static boolean chainContains(Throwable throwable, Class type) { + for (Throwable current = throwable; current != null; current = current.getCause()) { + if (type.isInstance(current)) { + return true; + } + } + return false; + } + private boolean shouldContinuePolling(TGetOperationStatusResp statusResp) { return statusResp == null || !statusResp.isSetOperationState() @@ -1042,10 +1295,35 @@ private TimeoutHandler getTimeoutHandler( TGetOperationStatusResp getOperationStatus( TGetOperationStatusReq statusReq, StatementId statementId) throws TException { long operationStatusStartTime = System.nanoTime(); - TGetOperationStatusResp operationStatus = getThriftClient().GetOperationStatus(statusReq); - long operationStatusEndTime = System.nanoTime(); - long operationStatusLatencyMillis = - (operationStatusEndTime - operationStatusStartTime) / 1_000_000; + TGetOperationStatusResp operationStatus = + withTransportRetry( + "GetOperationStatus", + statementId.toSQLExecStatementId(), + () -> getThriftClient().GetOperationStatus(statusReq)); + return recordOperationStatusLatency(statementId, operationStatusStartTime, operationStatus); + } + + /** + * Timeout-aware variant used by the execution polling loop: the retry backoff is bounded by the + * statement's {@code queryTimeout} via {@code timeoutHandler} so a transient transport failure + * cannot overshoot the deadline. + */ + TGetOperationStatusResp getOperationStatus( + TGetOperationStatusReq statusReq, StatementId statementId, TimeoutHandler timeoutHandler) + throws TException, DatabricksTimeoutException { + long operationStatusStartTime = System.nanoTime(); + TGetOperationStatusResp operationStatus = + withTransportRetry( + "GetOperationStatus", + statementId.toSQLExecStatementId(), + timeoutHandler, + () -> getThriftClient().GetOperationStatus(statusReq)); + return recordOperationStatusLatency(statementId, operationStatusStartTime, operationStatus); + } + + private TGetOperationStatusResp recordOperationStatusLatency( + StatementId statementId, long startTimeNanos, TGetOperationStatusResp operationStatus) { + long operationStatusLatencyMillis = (System.nanoTime() - startTimeNanos) / 1_000_000; LOGGER.debug( "Statement [{}] Thrift operation status latency: {}ms", statementId, diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java index 3bb04c8d6..a8127581d 100644 --- a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java @@ -5,20 +5,43 @@ /** Exception class to handle http errors while downloading chunk data from external links. */ public class DatabricksHttpException extends DatabricksSQLException { + /** + * HTTP status code that produced this exception, or {@code 0} when no HTTP response was received + * (e.g. a connection-level failure) or the status is otherwise unknown. + */ + private final int statusCode; + public DatabricksHttpException( String message, Throwable cause, DatabricksDriverErrorCode sqlCode) { super(message, cause, sqlCode); + this.statusCode = 0; } public DatabricksHttpException(String message, DatabricksDriverErrorCode internalCode) { super(message, null, internalCode.toString()); + this.statusCode = 0; } public DatabricksHttpException(String message, String sqlState) { super(message, null, sqlState); + this.statusCode = 0; } public DatabricksHttpException(String message, Throwable throwable, String sqlState) { super(message, throwable, sqlState); + this.statusCode = 0; + } + + public DatabricksHttpException(String message, String sqlState, int statusCode) { + super(message, null, sqlState); + this.statusCode = statusCode; + } + + /** + * Returns the HTTP status code associated with this exception, or {@code 0} when no HTTP response + * was received (connection-level failure) or the status is unknown. + */ + public int getStatusCode() { + return statusCode; } } diff --git a/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java index 384aa7761..2ac7313d0 100644 --- a/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java @@ -92,10 +92,11 @@ void testUnsuccessfulResponseCheck() { when(response.getStatusLine()).thenReturn(statusLine); when(statusLine.getStatusCode()).thenReturn(400); when(statusLine.toString()).thenReturn("mockStatusLine"); - Throwable exception = + DatabricksHttpException exception = assertThrows(DatabricksHttpException.class, () -> ValidationUtil.checkHTTPError(response)); assertEquals( "HTTP request failed by code: 400, status line: mockStatusLine.", exception.getMessage()); + assertEquals(400, exception.getStatusCode()); when(statusLine.getStatusCode()).thenReturn(102); assertThrows(DatabricksHttpException.class, () -> ValidationUtil.checkHTTPError(response)); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java index ae75478fc..a23a4ecd9 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java @@ -135,6 +135,38 @@ void testNullTimeoutAction() throws Exception { assertTrue(exception.getMessage().contains("timed-out after 2 seconds")); } + @Test + void testGetRemainingMillisStaysPositiveUntilEnforcement() throws Exception { + TimeoutHandler handler = + new TimeoutHandler( + 5, "Test operation", null, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + Field startTimeField = TimeoutHandler.class.getDeclaredField("startTimeMillis"); + startTimeField.setAccessible(true); + long now = System.currentTimeMillis(); + + // Elapsed 5.5s: past the nominal 5s but before checkTimeout enforces (it fires only once elapsed + // exceeds 5 whole seconds, i.e. at 6s). getRemainingMillis() must stay POSITIVE in this window + // so a caller capping a backoff sleep on it does not collapse the sleep to zero and busy-spin. + startTimeField.set(handler, now - 5_500L); + assertDoesNotThrow(handler::checkTimeout); + assertTrue( + handler.getRemainingMillis() > 0, + "remaining should be > 0 in the sub-second window before enforcement"); + + // Well past enforcement: remaining goes negative and checkTimeout throws. + startTimeField.set(handler, now - 10_000L); + assertTrue(handler.getRemainingMillis() < 0); + assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); + } + + @Test + void testGetRemainingMillisUnboundedWhenNoTimeout() { + TimeoutHandler handler = + new TimeoutHandler( + 0, "Test operation", null, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + assertEquals(Long.MAX_VALUE, handler.getRemainingMillis()); + } + @Test void testForStatementFactory() throws Exception { when(mockStatementId.toString()).thenReturn("test-statement-id"); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 58b65b5ba..25ea4fe11 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -14,14 +14,18 @@ import com.databricks.jdbc.common.StatementType; import com.databricks.jdbc.dbclient.impl.common.ClientConfigurator; import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.dbclient.impl.common.TimeoutHandler; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksRetryHandlerException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.exception.DatabricksValidationException; import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; +import java.net.SocketException; import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; @@ -1350,6 +1354,231 @@ void testExecute_remapsConcurrentModificationOnOperationStateBranchToSerializati assertEquals(1003, e.getErrorCode()); // EXECUTE_STATEMENT_FAILED stable code } + // --------------------------------------------------------------------------- + // Transport-retry classification and timeout behaviour (withTransportRetry). + // --------------------------------------------------------------------------- + + private static TTransportException transportError(Throwable cause) { + return new TTransportException(TTransportException.UNKNOWN, "transport failed", cause); + } + + @Test + void testTransportRetry_permanentHttpErrorSurfacesImmediately() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // A 404 (invalid handle) carries a concrete HTTP status: permanent, must not be retried. + TTransportException permanent = + transportError( + new DatabricksHttpException("HTTP request failed by code: 404", "08000", 404)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(permanent); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TTransportException thrown = + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + assertSame(permanent, thrown); + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_httpLayerHandledErrorNotReRetried() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 429/503 exhausted by the shared HTTP retry handler surface with a + // DatabricksRetryHandlerException in the cause chain — the outer loop must not pile on. + TTransportException exhausted = + transportError( + new DatabricksHttpException( + "Retry failure. HTTP response code: 429", + new DatabricksRetryHandlerException("rate limited", 429), + "08000")); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(exhausted); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_transientGatewayRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 502 is a transient gateway code the HTTP layer does not itself retry. + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(gateway) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + verify(accessor, times(1)).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_connectionFailureRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // A stale pooled connection / reset surfaces as a status-less DatabricksHttpException whose + // cause is a connection-level IOException (SocketException) — the PR's core case. + TTransportException connectionFailure = + transportError( + new DatabricksHttpException( + "Caught error while executing http request", + new SocketException("Connection reset"), + "08000")); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(connectionFailure) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + } + + @Test + void testTransportRetry_exhaustsThenRethrowsOriginal() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(gateway); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TTransportException thrown = + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + assertSame(gateway, thrown); + // 1 initial attempt + 5 bounded retries = 6 invocations, 5 backoff sleeps. + verify(thriftClient, times(6)).GetOperationStatus(operationStatusReq); + verify(accessor, times(5)).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_serverError500SurfacesImmediately() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 500 is treated as a (likely-permanent) server error, not a transient gateway hop failure. + TTransportException serverError = + transportError( + new DatabricksHttpException("HTTP request failed by code: 500", "08000", 500)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(serverError); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_nullCauseRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // DatabricksHttpTTransport.read() raises bare TTransportExceptions (null cause) on a + // truncated/empty response body — plausibly transient, so they are retried. + TTransportException bare = new TTransportException("Response buffer is empty, no response."); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(bare) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + } + + @Test + void testCleanupRetry_boundedForCancelOperation() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TCancelOperationReq req = + new TCancelOperationReq() + .setOperationHandle( + new TOperationHandle() + .setOperationId(handleIdentifier) + .setOperationType(TOperationType.UNKNOWN)); + // A persistent transient transport failure on cleanup must not run the full poll budget. + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.CancelOperation(req)).thenThrow(gateway); + + assertThrows(DatabricksHttpException.class, () -> accessor.cancelOperation(req)); + + // 1 initial attempt + MAX_CLEANUP_TRANSPORT_RETRIES (2) = 3 total, far below the poll budget. + verify(thriftClient, times(3)).CancelOperation(req); + verify(accessor, times(2)).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_consultsTimeoutHandlerBeforeSleeping() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TimeoutHandler timeoutHandler = mock(TimeoutHandler.class); + when(timeoutHandler.getRemainingMillis()).thenReturn(10_000L); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(gateway) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = + accessor.getOperationStatus(operationStatusReq, statementId, timeoutHandler); + + assertSame(operationStatusFinishedResp, resp); + verify(timeoutHandler, atLeastOnce()).checkTimeout(); + } + + @Test + void testTransportRetry_timeoutAbortsRetryWithinDeadline() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TimeoutHandler timeoutHandler = mock(TimeoutHandler.class); + doThrow( + new DatabricksTimeoutException( + "Statement execution timed-out", + /* cause= */ null, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT)) + .when(timeoutHandler) + .checkTimeout(); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(gateway); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + DatabricksTimeoutException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId, timeoutHandler)); + + // The first attempt fails (retryable), the deadline fires before any backoff sleep, so no + // further RPC attempts and no sleep — the retry cannot overshoot queryTimeout. + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + private TFetchResultsReq getFetchResultsRequest(boolean includeMetadata) throws DatabricksValidationException { TFetchResultsReq request =