Skip to content

Commit c13e298

Browse files
sreekanth-dbsreekanth-db
authored
SQLTimeoutException fix (#1135)
## Description <!-- Provide a brief summary of the changes made and the issue they aim to address.--> This PR fixes timeout exception handling to ensure `SQLTimeoutException` is thrown (via `DatabricksTimeoutException`) instead of generic `DatabricksSQLException` when queries timeout, improving JDBC standards compliance. **Changes:** - **Client-side timeout fix** (`DatabricksStatement.java`): Changed exception unwrapping to check for `SQLException` instead of just `DatabricksSQLException`, allowing `DatabricksTimeoutException` to propagate correctly - **Server-side timeout fix for Thrift mode** (`DatabricksThriftAccessor.java`): Added SQL state check for `57KD0` (Databricks timeout code) in `checkOperationStatusForErrors()` to throw `DatabricksTimeoutException` - **Server-side timeout fix for SEA mode** (`DatabricksSdkClient.java`): Added SQL state check for `57KD0` in `handleFailedExecution()` to throw `DatabricksTimeoutException` - **Unit tests**: Added server-side timeout tests for both Thrift and SEA modes ## Testing <!-- Describe how the changes have been tested--> - Existing client-side timeout tests pass (`testClientSideTimeout` in both test files) - New server-side timeout tests pass (`testServerSideTimeoutThrowsTimeoutException` in both test files) - Manually verified expected exceptions by connecting to a real Databricks workspace ## Additional Notes to the Reviewer <!-- Share any additional context or insights that may help the reviewer understand the changes better. This could include challenges faced, limitations, or compromises made during the development process. Also, mention any areas of the code that you would like the reviewer to focus on specifically. --> --------- Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
1 parent 6e1c164 commit c13e298

8 files changed

Lines changed: 126 additions & 27 deletions

File tree

‎NEXT_CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,5 +8,7 @@
88

99
### Fixed
1010

11+
- Fix timeout exception handling to throw `SQLTimeoutException` instead of `DatabricksSQLException` when queries timeout.
12+
1113
---
1214
*Note: When making changes, please add your change under the appropriate section with a brief description.*

‎development/.release-freeze.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
2-
"freeze": true,
3-
"reason": "Release freeze for v3.0.6"
2+
"freeze": false,
3+
"reason": ""
44
}

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -755,11 +755,11 @@ DatabricksResultSet executeInternal(
755755
timeoutErrorMessage, e, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT);
756756
} catch (InterruptedException | ExecutionException e) {
757757
Throwable cause = e;
758-
// Look for underlying DatabricksSQL exception
758+
// Look for underlying SQLException (includes DatabricksSQLException and other SQL exceptions)
759759
while (cause.getCause() != null) {
760760
cause = cause.getCause();
761-
if (cause instanceof DatabricksSQLException) {
762-
throw (DatabricksSQLException) cause;
761+
if (cause instanceof SQLException) {
762+
throw (SQLException) cause;
763763
}
764764
}
765765
String errMsg =

‎src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ public final class DatabricksJdbcConstants {
104104
public static final String GCP_GOOGLE_CREDENTIALS_AUTH_TYPE = "google-credentials";
105105
public static final String GCP_GOOGLE_ID_AUTH_TYPE = "google-id";
106106
public static final String DEFAULT_HTTP_EXCEPTION_SQLSTATE = "08000";
107+
public static final String QUERY_EXECUTION_TIMEOUT_SQLSTATE = "57KD0";
107108
public static final int TEMPORARY_REDIRECT_STATUS_CODE = 307;
108109
public static final String REDACTED_TOKEN = "****";
109110
public static final int MAX_DEFAULT_STRING_COLUMN_LENGTH = 32767;

‎src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.databricks.jdbc.dbclient.impl.sqlexec;
22

33
import static com.databricks.jdbc.common.DatabricksJdbcConstants.JSON_HTTP_HEADERS;
4+
import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE;
45
import static com.databricks.jdbc.common.DatabricksJdbcConstants.TEMPORARY_REDIRECT_STATUS_CODE;
56
import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT;
67
import static com.databricks.jdbc.common.util.DatabricksTypeUtil.DECIMAL;
@@ -650,10 +651,14 @@ void handleFailedExecution(
650651
" Error Message: %s, Error code: %s", error.getMessage(), error.getErrorCode());
651652
}
652653
LOGGER.debug(errorMessage);
654+
655+
String sqlState = response.getStatus().getSqlState();
656+
if (QUERY_EXECUTION_TIMEOUT_SQLSTATE.equals(sqlState)) {
657+
throw new DatabricksTimeoutException(
658+
errorMessage, null, DatabricksDriverErrorCode.OPERATION_TIMEOUT_ERROR);
659+
}
653660
throw new DatabricksSQLException(
654-
errorMessage,
655-
response.getStatus().getSqlState(),
656-
DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED);
661+
errorMessage, sqlState, DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED);
657662
}
658663

659664
private ExecuteStatementResponse wrapGetStatementResponse(

‎src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java‎

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.databricks.jdbc.dbclient.impl.thrift;
22

3+
import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE;
34
import static com.databricks.jdbc.common.EnvironmentVariables.*;
45
import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*;
56

@@ -574,63 +575,58 @@ TFetchResultsResp fetchResultsWithAbsoluteOffset(
574575
}
575576

576577
private TFetchResultsResp listFunctions(TGetFunctionsReq request)
577-
throws TException, DatabricksSQLException {
578+
throws TException, SQLException {
578579
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
579580
TGetFunctionsResp response = getThriftClient().GetFunctions(request);
580581
return fetchMetadataResults(response, response.toString());
581582
}
582583

583584
private TFetchResultsResp listPrimaryKeys(TGetPrimaryKeysReq request)
584-
throws TException, DatabricksSQLException {
585+
throws TException, SQLException {
585586
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
586587
TGetPrimaryKeysResp response = getThriftClient().GetPrimaryKeys(request);
587588
return fetchMetadataResults(response, response.toString());
588589
}
589590

590591
private TFetchResultsResp listCrossReferences(TGetCrossReferenceReq request)
591-
throws TException, DatabricksSQLException {
592+
throws TException, SQLException {
592593
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
593594
TGetCrossReferenceResp response = getThriftClient().GetCrossReference(request);
594595
return fetchMetadataResults(response, response.toString());
595596
}
596597

597-
private TFetchResultsResp getTables(TGetTablesReq request)
598-
throws TException, DatabricksSQLException {
598+
private TFetchResultsResp getTables(TGetTablesReq request) throws TException, SQLException {
599599
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
600600
TGetTablesResp response = getThriftClient().GetTables(request);
601601
return fetchMetadataResults(response, response.toString());
602602
}
603603

604604
private TFetchResultsResp getTableTypes(TGetTableTypesReq request)
605-
throws TException, DatabricksSQLException {
605+
throws TException, SQLException {
606606
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
607607
TGetTableTypesResp response = getThriftClient().GetTableTypes(request);
608608
return fetchMetadataResults(response, response.toString());
609609
}
610610

611-
private TFetchResultsResp getCatalogs(TGetCatalogsReq request)
612-
throws TException, DatabricksSQLException {
611+
private TFetchResultsResp getCatalogs(TGetCatalogsReq request) throws TException, SQLException {
613612
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
614613
TGetCatalogsResp response = getThriftClient().GetCatalogs(request);
615614
return fetchMetadataResults(response, response.toString());
616615
}
617616

618-
private TFetchResultsResp listSchemas(TGetSchemasReq request)
619-
throws TException, DatabricksSQLException {
617+
private TFetchResultsResp listSchemas(TGetSchemasReq request) throws TException, SQLException {
620618
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
621619
TGetSchemasResp response = getThriftClient().GetSchemas(request);
622620
return fetchMetadataResults(response, response.toString());
623621
}
624622

625-
private TFetchResultsResp getTypeInfo(TGetTypeInfoReq request)
626-
throws TException, DatabricksSQLException {
623+
private TFetchResultsResp getTypeInfo(TGetTypeInfoReq request) throws TException, SQLException {
627624
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
628625
TGetTypeInfoResp response = getThriftClient().GetTypeInfo(request);
629626
return fetchMetadataResults(response, response.toString());
630627
}
631628

632-
private TFetchResultsResp listColumns(TGetColumnsReq request)
633-
throws TException, DatabricksSQLException {
629+
private TFetchResultsResp listColumns(TGetColumnsReq request) throws TException, SQLException {
634630
if (enableDirectResults) request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
635631
TGetColumnsResp response = getThriftClient().GetColumns(request);
636632
return fetchMetadataResults(response, response.toString());
@@ -669,7 +665,7 @@ private TCLIService.Client newThriftClient() {
669665
*/
670666
private <TResp extends TBase<TResp, FResp>, FResp extends TFieldIdEnum>
671667
TFetchResultsResp fetchMetadataResults(TResp response, String contextDescription)
672-
throws TException, DatabricksSQLException {
668+
throws TException, SQLException {
673669
checkResponseForErrors(response);
674670

675671
// Get the operation status from direct results if present
@@ -742,7 +738,7 @@ private <T extends TBase<T, F>, F extends TFieldIdEnum> void checkResponseForErr
742738
}
743739

744740
private void checkOperationStatusForErrors(TGetOperationStatusResp statusResp, String statementId)
745-
throws DatabricksSQLException {
741+
throws SQLException {
746742
if (statusResp != null
747743
&& statusResp.isSetOperationState()
748744
&& isErrorOperationState(statusResp.getOperationState())) {
@@ -751,7 +747,14 @@ && isErrorOperationState(statusResp.getOperationState())) {
751747
"Operation failed with error: [%s] for statement [%s], with response [%s]",
752748
statusResp.getErrorMessage(), statementId, statusResp);
753749
LOGGER.error(errorMsg);
754-
throw new DatabricksSQLException(errorMsg, statusResp.getSqlState());
750+
751+
String sqlState = statusResp.getSqlState();
752+
if (QUERY_EXECUTION_TIMEOUT_SQLSTATE.equals(sqlState)) {
753+
throw new DatabricksTimeoutException(
754+
errorMsg, null, DatabricksDriverErrorCode.OPERATION_TIMEOUT_ERROR);
755+
}
756+
757+
throw new DatabricksSQLException(errorMsg, sqlState);
755758
}
756759
}
757760

‎src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.databricks.jdbc.dbclient.impl.sqlexec;
22

33
import static com.databricks.jdbc.TestConstants.TEST_STRING;
4+
import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE;
45
import static com.databricks.jdbc.common.DatabricksJdbcConstants.TEMPORARY_REDIRECT_STATUS_CODE;
56
import static com.databricks.jdbc.dbclient.impl.sqlexec.PathConstants.*;
67
import static com.databricks.jdbc.model.core.ColumnInfoTypeName.DECIMAL;
@@ -438,6 +439,59 @@ public void testExecuteStatementWithTimeoutExpired() throws Exception {
438439
verify(databricksSdkClient).cancelStatement(eq(STATEMENT_ID));
439440
}
440441

442+
@Test
443+
public void testServerSideTimeoutThrowsTimeoutException() throws Exception {
444+
445+
IDatabricksConnectionContext connectionContext =
446+
DatabricksConnectionContext.parse(JDBC_URL, new Properties());
447+
DatabricksSdkClient databricksSdkClient =
448+
spy(new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient));
449+
DatabricksConnection connection =
450+
new DatabricksConnection(connectionContext, databricksSdkClient);
451+
452+
// Mock session creation
453+
CreateSessionResponse sessionResponse = new CreateSessionResponse().setSessionId(SESSION_ID);
454+
when(apiClient.execute(any(Request.class), eq(CreateSessionResponse.class)))
455+
.thenReturn(sessionResponse);
456+
connection.open();
457+
458+
// Create statement with a long client timeout (server times out first)
459+
DatabricksStatement statement = new DatabricksStatement(connection);
460+
statement.setMaxRows(100);
461+
statement.setQueryTimeout(300); // 300 seconds - server will timeout before this
462+
463+
// Mock server-side timeout: server returns FAILED state with sqlState="57KD0"
464+
ExecuteStatementResponse executeResponse =
465+
new ExecuteStatementResponse()
466+
.setStatementId(STATEMENT_ID.toSQLExecStatementId())
467+
.setStatus(
468+
new StatementStatus()
469+
.setState(StatementState.FAILED)
470+
.setSqlState(QUERY_EXECUTION_TIMEOUT_SQLSTATE) // Server-side timeout SQL state
471+
.setError(
472+
new ServiceError()
473+
.setMessage("Statement has timed out after 10 seconds.")
474+
.setErrorCode(ServiceErrorCode.BAD_REQUEST)));
475+
476+
// Set up mock response
477+
when(apiClient.execute(
478+
argThat(req -> req != null && STATEMENT_PATH.equals(req.getUrl())),
479+
eq(ExecuteStatementResponse.class)))
480+
.thenReturn(executeResponse);
481+
482+
// Verify that DatabricksTimeoutException is thrown
483+
assertThrows(
484+
DatabricksTimeoutException.class,
485+
() ->
486+
databricksSdkClient.executeStatement(
487+
STATEMENT,
488+
warehouse,
489+
sqlParams,
490+
StatementType.QUERY,
491+
connection.getSession(),
492+
statement));
493+
}
494+
441495
@Test
442496
public void testDecimalTypeWithValidPrecisionAndScale() throws DatabricksSQLException {
443497
BigDecimal decimalValue = new BigDecimal("123.45"); // precision: 5, scale: 2

‎src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.databricks.jdbc.dbclient.impl.thrift;
22

3+
import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE;
34
import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_BYTE_LIMIT;
45
import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK;
56
import static org.junit.jupiter.api.Assertions.*;
@@ -824,8 +825,7 @@ void testExecuteWithTimeout() throws TException, SQLException, DatabricksValidat
824825
}
825826

826827
@Test
827-
void testExecuteWithTimeoutExpired()
828-
throws TException, SQLException, DatabricksValidationException {
828+
void testExecuteWithTimeoutExpired() throws TException, SQLException {
829829
// Set the async poll interval to 1 second to facilitate testing
830830
when(connectionContext.getAsyncExecPollInterval()).thenReturn(1000);
831831

@@ -869,6 +869,40 @@ void testExecuteWithTimeoutExpired()
869869
verify(thriftClient).CancelOperation(any(TCancelOperationReq.class));
870870
}
871871

872+
@Test
873+
void testServerSideTimeoutThrowsTimeoutException() throws TException, SQLException {
874+
875+
accessor = spy(new DatabricksThriftAccessor(connectionContext));
876+
doReturn(thriftClient).when(accessor).getThriftClient();
877+
878+
// Create statement execution mocks
879+
TExecuteStatementReq request = new TExecuteStatementReq();
880+
TExecuteStatementResp tExecuteStatementResp =
881+
new TExecuteStatementResp()
882+
.setOperationHandle(tOperationHandle)
883+
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
884+
when(thriftClient.ExecuteStatement(request)).thenReturn(tExecuteStatementResp);
885+
886+
// Mock server-side timeout: server returns ERROR_STATE with sqlState="57KD0"
887+
TGetOperationStatusResp operationStatusErrorResp =
888+
new TGetOperationStatusResp()
889+
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
890+
.setOperationState(TOperationState.ERROR_STATE)
891+
.setSqlState(QUERY_EXECUTION_TIMEOUT_SQLSTATE) // Server-side timeout SQL state
892+
.setErrorMessage("Statement has timed out after 10 seconds.");
893+
894+
when(thriftClient.GetOperationStatus(operationStatusReq)).thenReturn(operationStatusErrorResp);
895+
896+
Statement statement = mock(Statement.class);
897+
when(parentStatement.getStatement()).thenReturn(statement);
898+
when(statement.getQueryTimeout()).thenReturn(300); // Long timeout, server times out first
899+
900+
// Verify that DatabricksTimeoutException is thrown
901+
assertThrows(
902+
DatabricksTimeoutException.class,
903+
() -> accessor.execute(request, parentStatement, session, StatementType.SQL));
904+
}
905+
872906
@Test
873907
void testFetchResultsWithCustomMaxRowsPerBlock()
874908
throws TException, SQLException, DatabricksValidationException {

0 commit comments

Comments
 (0)