[ES-1717770] Fix TIMEDOUT_STATE not recognized as error on interactive clusters

NEXT_CHANGELOG.md

@@ -13,6 +13,7 @@
 - PECOBLR-1121 Arrow patch to circumvent Arrow issues with JDK 16+.
 
 ### Fixed
+- Fixed statement timeout when the server returns `TIMEDOUT_STATE` directly in the `ExecuteStatement` response (e.g. query queued under load), the driver now throws `SQLTimeoutException` instead of `DatabricksHttpException`.
 - Fixed Thrift polling infinite loop when server restarts invalidate operation handles, and added configurable timeout (`MetadataOperationTimeout`, default 300s) with sleep between polls for metadata operations.
 - Fixed `DatabricksParameterMetaData.countParameters` and `DatabricksStatement.trimCommentsAndWhitespaces` with a `SqlCommentParser` utility class.
 - Fixed `rollback()` to throw `SQLException` when called in auto-commit mode (no active transaction), aligning with JDBC spec. Previously it silently sent a ROLLBACK command to the server.

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

@@ -817,7 +817,8 @@ private void checkOperationStatusForErrors(TGetOperationStatusResp statusResp, S
       LOGGER.error(errorMsg);
 
       String sqlState = statusResp.getSqlState();
-      if (QUERY_EXECUTION_TIMEOUT_SQLSTATE.equals(sqlState)) {
+      if (QUERY_EXECUTION_TIMEOUT_SQLSTATE.equals(sqlState)
+          || statusResp.getOperationState() == TOperationState.TIMEDOUT_STATE) {
         throw new DatabricksTimeoutException(
             errorMsg, null, DatabricksDriverErrorCode.OPERATION_TIMEOUT_ERROR);
       }
@@ -854,7 +855,9 @@ private boolean isErrorStatusCode(TStatus status) {
   }
 
   private boolean isErrorOperationState(TOperationState state) {
-    return state == TOperationState.ERROR_STATE || state == TOperationState.CLOSED_STATE;
+    return state == TOperationState.ERROR_STATE
+        || state == TOperationState.CLOSED_STATE
+        || state == TOperationState.TIMEDOUT_STATE;
   }
 
   private boolean isPendingOperationState(TOperationState state) {

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

@@ -893,6 +893,70 @@ void testServerSideTimeoutThrowsTimeoutException() throws TException, SQLExcepti
         () -> accessor.execute(request, parentStatement, session, StatementType.SQL));
   }
 
+  @Test
+  void testTimedOutStateInDirectResultsThrowsTimeoutException()
+      throws TException, SQLException, DatabricksValidationException {
+    // Reproduces the interactive cluster scenario: server enforces queryTimeout and returns
+    // TIMEDOUT_STATE directly in directResults before the client polling loop starts (e.g. query
+    // is queued under load and times out while waiting). Previously isErrorOperationState excluded
+    // TIMEDOUT_STATE, causing the driver to fall through to executeFetchRequest and throw
+    // DatabricksHttpException instead.
+    setup(true);
+
+    TExecuteStatementReq request = new TExecuteStatementReq();
+    TSparkDirectResults timedOutDirectResults =
+        new TSparkDirectResults()
+            .setOperationStatus(
+                new TGetOperationStatusResp()
+                    .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
+                    .setOperationState(TOperationState.TIMEDOUT_STATE)
+                    .setErrorMessage("Query timed out after 1 seconds"));
+    TExecuteStatementResp tExecuteStatementResp =
+        new TExecuteStatementResp()
+            .setOperationHandle(tOperationHandle)
+            .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
+            .setDirectResults(timedOutDirectResults);
+    when(thriftClient.ExecuteStatement(request)).thenReturn(tExecuteStatementResp);
+
+    Statement statement = mock(Statement.class);
+    when(parentStatement.getStatement()).thenReturn(statement);
+    when(statement.getQueryTimeout()).thenReturn(300); // Long client timeout — server fires first
+
+    assertThrows(
+        DatabricksTimeoutException.class,
+        () -> accessor.execute(request, parentStatement, session, StatementType.SQL));
+  }
+
+  @Test
+  void testTimedOutStateDuringPollingThrowsTimeoutException()
+      throws TException, SQLException, DatabricksValidationException {
+    // Server returns RUNNING_STATE initially, then TIMEDOUT_STATE during polling —
+    // e.g. cluster enforces its own max query duration while client timeout is longer.
+    setup(true);
+
+    TExecuteStatementReq request = new TExecuteStatementReq();
+    TExecuteStatementResp tExecuteStatementResp =
+        new TExecuteStatementResp()
+            .setOperationHandle(tOperationHandle)
+            .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
+    when(thriftClient.ExecuteStatement(request)).thenReturn(tExecuteStatementResp);
+
+    TGetOperationStatusResp timedOutStatusResp =
+        new TGetOperationStatusResp()
+            .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
+            .setOperationState(TOperationState.TIMEDOUT_STATE)
+            .setErrorMessage("Query timed out after 1 seconds");
+    when(thriftClient.GetOperationStatus(operationStatusReq)).thenReturn(timedOutStatusResp);
+
+    Statement statement = mock(Statement.class);
+    when(parentStatement.getStatement()).thenReturn(statement);
+    when(statement.getQueryTimeout()).thenReturn(300); // Long client timeout — server fires first
+
+    assertThrows(
+        DatabricksTimeoutException.class,
+        () -> accessor.execute(request, parentStatement, session, StatementType.SQL));
+  }
+
   @Test
   void testFetchResultsWithCustomMaxRowsPerBlock()
       throws TException, SQLException, DatabricksValidationException {