From ce39904d212898e361c40a1d61975f4d6b83e84a Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 23 Sep 2026 00:27:03 +0000 Subject: [PATCH] fix: return HYT00 for query timeouts Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- NEXT_CHANGELOG.md | 2 ++ .../jdbc/api/impl/DatabricksStatement.java | 5 +++- .../jdbc/common/DatabricksJdbcConstants.java | 3 +++ .../dbclient/impl/common/TimeoutHandler.java | 5 +++- .../api/impl/DatabricksStatementTest.java | 27 +++++++++++++++++++ .../impl/common/TimeoutHandlerTest.java | 2 ++ .../impl/sqlexec/DatabricksSdkClientTest.java | 2 ++ .../thrift/DatabricksThriftAccessorTest.java | 2 ++ 8 files changed, 46 insertions(+), 2 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index d27a25e263..2cdc7d5cff 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,8 @@ ### Fixed +- Return SQLSTATE `HYT00` when a client-side query timeout expires. + --- *Note: When making changes, please add your change under the appropriate section with a brief description.* diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index d1dd0d30e1..b33f323e4a 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -890,7 +890,10 @@ DatabricksResultSet executeInternal( LOGGER.error(timeoutErrorMessage); futureResultSet.cancel(true); // Cancel execution run throw new DatabricksTimeoutException( - timeoutErrorMessage, e, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + timeoutErrorMessage, + TIMEOUT_EXPIRED_SQLSTATE, + e, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); } catch (InterruptedException | ExecutionException e) { Throwable cause = e; // Look for underlying SQLException (includes DatabricksSQLException and other SQL exceptions) diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java index 3e874a2566..0dcdb9ebf0 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java @@ -114,6 +114,9 @@ public final class DatabricksJdbcConstants { public static final String DEFAULT_HTTP_EXCEPTION_SQLSTATE = OPERATION_ERROR_SQLSTATE; public static final String QUERY_EXECUTION_TIMEOUT_SQLSTATE = "57KD0"; + /** SQLSTATE for timeout expiration, matching the Databricks reference JDBC driver. */ + public static final String TIMEOUT_EXPIRED_SQLSTATE = "HYT00"; + /** Standard SQL state for syntax error or access rule violation (SQLSTATE 42000). */ public static final String SYNTAX_OR_ACCESS_VIOLATION_SQLSTATE = "42000"; 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 edc85e454d..7756e839cf 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 @@ -1,5 +1,7 @@ package com.databricks.jdbc.dbclient.impl.common; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.TIMEOUT_EXPIRED_SQLSTATE; + import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.log.JdbcLogger; @@ -68,7 +70,8 @@ public void checkTimeout() throws DatabricksTimeoutException { "Statement execution timed-out after %d seconds. Operation: %s", timeoutSeconds, operationDescription); LOGGER.error(timeoutErrorMessage); - throw new DatabricksTimeoutException(timeoutErrorMessage, /*cause*/ null, internalErrorCode); + throw new DatabricksTimeoutException( + timeoutErrorMessage, TIMEOUT_EXPIRED_SQLSTATE, /*cause*/ null, internalErrorCode); } } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index e1dc979841..7b91ea7fb2 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -1,5 +1,6 @@ package com.databricks.jdbc.api.impl; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.TIMEOUT_EXPIRED_SQLSTATE; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -15,6 +16,7 @@ import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksSdkClient; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException; +import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.model.core.StatementStatus; import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.service.sql.StatementState; @@ -330,6 +332,31 @@ public void testExecuteInternalWithZeroTimeout() throws Exception { verify(mockFuture, never()).get(anyLong(), any(TimeUnit.class)); } + @Test + public void testExecuteInternalTimeoutUsesTimeoutExpiredSqlState() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection mockConnection = mock(DatabricksConnection.class); + when(mockConnection.getConnectionContext()).thenReturn(connectionContext); + DatabricksStatement statement = spy(new DatabricksStatement(mockConnection)); + statement.setQueryTimeout(1); + + CompletableFuture mockFuture = mock(CompletableFuture.class); + when(mockFuture.get(1, TimeUnit.SECONDS)) + .thenThrow(new java.util.concurrent.TimeoutException()); + doReturn(mockFuture).when(statement).getFutureResult(anyString(), anyMap(), any()); + + DatabricksTimeoutException exception = + assertThrows( + DatabricksTimeoutException.class, + () -> + statement.executeInternal( + "SELECT * FROM table", new HashMap<>(), StatementType.QUERY, false)); + + assertEquals(TIMEOUT_EXPIRED_SQLSTATE, exception.getSQLState()); + verify(mockFuture).cancel(true); + } + @Test public void testInputStreamForVolumeOperation() throws Exception { IDatabricksConnectionContext connectionContext = 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 ae75478fc3..0f9aa187c0 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 @@ -1,5 +1,6 @@ package com.databricks.jdbc.dbclient.impl.common; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.TIMEOUT_EXPIRED_SQLSTATE; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; @@ -82,6 +83,7 @@ void testTimeoutActionExecuted() throws Exception { // Verify exception message assertTrue(exception.getMessage().contains("timed-out after 2 seconds")); assertTrue(exception.getMessage().contains("Test operation")); + assertEquals(TIMEOUT_EXPIRED_SQLSTATE, exception.getSQLState()); } @Test diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index fc4a2efc0d..b12e844120 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -3,6 +3,7 @@ import static com.databricks.jdbc.TestConstants.TEST_STRING; import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.DatabricksJdbcConstants.TEMPORARY_REDIRECT_STATUS_CODE; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.TIMEOUT_EXPIRED_SQLSTATE; import static com.databricks.jdbc.dbclient.impl.sqlexec.PathConstants.*; import static com.databricks.jdbc.model.core.ColumnInfoTypeName.DECIMAL; import static com.databricks.jdbc.model.core.ColumnInfoTypeName.INT; @@ -782,6 +783,7 @@ public void testExecuteStatementWithTimeoutExpired() throws Exception { null)); assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); + assertEquals(TIMEOUT_EXPIRED_SQLSTATE, exception.getSQLState()); // Verify cancel was called verify(databricksSdkClient).cancelStatement(eq(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 7f58270c15..0bbbc3ace8 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 @@ -1,6 +1,7 @@ package com.databricks.jdbc.dbclient.impl.thrift; import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.TIMEOUT_EXPIRED_SQLSTATE; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_BYTE_LIMIT; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK; import static com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED; @@ -954,6 +955,7 @@ void testExecuteWithTimeoutExpired() throws TException, SQLException { () -> accessor.execute(request, parentStatement, session, StatementType.SQL)); assertTrue(exception.getMessage().contains("timed-out after 1 seconds")); + assertEquals(TIMEOUT_EXPIRED_SQLSTATE, exception.getSQLState()); // Verify that cancel was called verify(thriftClient).CancelOperation(any(TCancelOperationReq.class));