Fix SEA result expiry error handling — preserve server error message (#1302)

## Summary
When query results expire after the retention period, the SEA server
returns HTTP 404 with a descriptive `NOT_FOUND: The result of the
statement <id> was not found.` message. The SDK client throws
`DatabricksError` (extends `RuntimeException`) which was not caught by
`getResultChunks`/`getResultChunksData` — they only caught
`IOException`.

This caused three problems:
1. **Server error message lost**: `DatabricksError` escaped uncaught,
and the `ChunkDownloadTask` finally block created a new exception
without chaining the original cause — customer sees `"Download failed
for chunk index N"` with no indication of why
2. **Download tasks hang indefinitely**: In `ChunkLinkDownloadService`,
the uncaught `RuntimeException` left link futures unresolved, blocking
download tasks forever on `future.get()`
3. **Inconsistent with other methods**: `createSession` already catches
`DatabricksError` properly — `getResultChunks`/`getResultChunksData`
should follow the same pattern

## Changes
- **`DatabricksSdkClient.java`**: Add `catch (DatabricksError e)` in
`getResultChunks()` and `getResultChunksData()`, wrapping into
`DatabricksSQLException` with HTTP status code and server message
preserved
- **`ChunkLinkDownloadService.java`**: Widen catch from `SQLException`
to `Exception` so `DatabricksError` completes link futures exceptionally
instead of leaving them unresolved
- **`ChunkDownloadTask.java`**: Chain `uncaughtException` as cause in
the finally block so the original error propagates to the caller
- **`ChunkLinkDownloadServiceTest.java`**: Fix `testHandleExpiredLinks`
to mock subsequent chunks needed by batch continuation

## What customers see after this fix
- **Before**: `SQLException: "Download failed for chunk index N"` (no
server message, potential hang)
- **After**: `SQLException: "Error fetching result chunks for statement
[id] chunk [N] (HTTP 404): NOT_FOUND: The result of the statement <id>
was not found."` (with full cause chain)

## Test plan
- [x] Full jdbc-core suite: 3170 tests pass, 0 failures
- [x] `ChunkLinkDownloadServiceTest` — 8 tests pass (including fixed
`testHandleExpiredLinks`)

This pull request was AI-assisted by Isaac.

---------

Signed-off-by: Gopal Lal <gopal.lal@databricks.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -18,6 +18,7 @@
 - Fixed primitive types within complex types (ARRAY, MAP, STRUCT) not being correctly parsed when Arrow serialization uses alternate formats: TIMESTAMP/TIMESTAMP_NTZ as epoch microseconds or component arrays, and BINARY as base64-encoded strings.
 - Fixed `PARSE_SYNTAX_ERROR` for column names containing special characters (e.g., dots) when `EnableBatchedInserts` is enabled, by re-quoting column names with backticks in reconstructed multi-row INSERT statements.
 - Fixed Volume ingestion for SEA mode, which was broken due to statement being closed prematurely.
+- Fixed preserving error message during query results expiry and graceful handling of RuntimeException.
 - Fixed escaped pattern characters in catalogName for `getSchemas`, as returned catalogName should be unescaped.
 - Fixed `getColumnClassName()` returning null for VARIANT columns in SEA mode by adding VARIANT to the type system.
 - Fixed `getColumns()` returning `DATA_TYPE=0` (NULL) for GEOMETRY/GEOGRAPHY columns in Thrift mode. Now returns `Types.VARCHAR` (12) when geospatial is disabled and `Types.OTHER` (1111) when enabled, consistent with SEA mode.

src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java

@@ -133,6 +133,7 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
             .completeExceptionally(
                 new DatabricksSQLException(
                     "Download failed for chunk index " + chunk.getChunkIndex(),
+                    uncaughtException,
                     DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR));
       }
 

src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java

@@ -11,6 +11,7 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.core.ChunkLinkFetchResult;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.sdk.core.DatabricksError;
 import com.google.common.annotations.VisibleForTesting;
 import java.sql.SQLException;
 import java.time.Instant;
@@ -286,8 +287,10 @@ private void triggerNextBatchDownload() {
                     triggerNextBatchDownload();
                   }
                 }
-              } catch (SQLException e) {
-                // If the download fails, complete exceptionally all pending futures
+              } catch (SQLException | DatabricksError e) {
+                // DatabricksError (RuntimeException) is thrown by the SDK client on HTTP
+                // errors (e.g. 404 when result expires). Without catching it, link futures
+                // are left unresolved, causing download tasks to block indefinitely.
                 handleBatchDownloadError(batchStartIndex, e);
               }
             });
@@ -299,7 +302,7 @@ private void triggerNextBatchDownload() {
    * <p>Completes all pending futures exceptionally with the encountered error and resets the
    * download progress flag.
    */
-  private void handleBatchDownloadError(long batchStartIndex, SQLException e) {
+  private void handleBatchDownloadError(long batchStartIndex, Exception e) {
     LOGGER.error(
         e,
         "Failed to download links for batch starting at {} : {}",

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

@@ -488,6 +488,13 @@ public ChunkLinkFetchResult getResultChunks(
       req.withHeaders(getHeaders("getStatementResultN"));
       ResultData resultData = apiClient.execute(req, ResultData.class);
       return buildChunkLinkFetchResult(resultData.getExternalLinks());
+    } catch (DatabricksError e) {
+      String errorMessage =
+          String.format(
+              "Error fetching result chunks for statement [%s] chunk [%d] (HTTP %d): %s",
+              statementId, chunkIndex, e.getStatusCode(), e.getMessage());
+      LOGGER.error(errorMessage, e);
+      throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.SDK_CLIENT_ERROR);
     } catch (IOException e) {
       String errorMessage = "Error while processing the get result chunk request";
       LOGGER.error(errorMessage, e);
@@ -541,6 +548,13 @@ public ResultData getResultChunksData(StatementId typedStatementId, long chunkIn
       Request req = new Request(Request.GET, path, apiClient.serialize(request));
       req.withHeaders(getHeaders("getStatementResultN"));
       return apiClient.execute(req, ResultData.class);
+    } catch (DatabricksError e) {
+      String errorMessage =
+          String.format(
+              "Error fetching result data for statement [%s] chunk [%d] (HTTP %d): %s",
+              statementId, chunkIndex, e.getStatusCode(), e.getMessage());
+      LOGGER.error(errorMessage, e);
+      throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.SDK_CLIENT_ERROR);
     } catch (IOException e) {
       String errorMessage = "Error while processing the get result chunk request";
       LOGGER.error(errorMessage, e);

src/test/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadServiceTest.java

@@ -14,6 +14,7 @@
 import com.databricks.jdbc.model.core.ChunkLinkFetchResult;
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import com.databricks.sdk.core.DatabricksError;
 import java.sql.SQLException;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
@@ -127,6 +128,28 @@ void testGetLinkForChunk_ClientError()
     assertEquals(expectedError, exception.getCause());
   }
 
+  @Test
+  void testGetLinkForChunk_DatabricksError_propagatesAsExecutionException()
+      throws SQLException, ExecutionException, InterruptedException {
+    long chunkIndex = 1L;
+    // DatabricksError is a RuntimeException thrown by the SDK on HTTP errors (e.g. 404)
+    DatabricksError runtimeError = new DatabricksError("404", "Results have expired", 404);
+    when(mockSession.getDatabricksClient()).thenReturn(mockClient);
+    when(mockClient.getResultChunks(eq(mockStatementId), anyLong(), anyLong()))
+        .thenThrow(runtimeError);
+    when(mockChunkMap.get(chunkIndex)).thenReturn(mock(ArrowResultChunk.class));
+
+    ChunkLinkDownloadService<ArrowResultChunk> service =
+        new ChunkLinkDownloadService<>(
+            mockSession, mockStatementId, TOTAL_CHUNKS, mockChunkMap, NEXT_BATCH_START_INDEX);
+
+    CompletableFuture<ExternalLink> future = service.getLinkForChunk(chunkIndex);
+
+    ExecutionException exception =
+        assertThrows(ExecutionException.class, () -> future.get(1, TimeUnit.SECONDS));
+    assertEquals(runtimeError, exception.getCause());
+  }
+
   @Test
   void testAutoTriggerForSEAClient() throws SQLException, InterruptedException {
     when(mockSession.getDatabricksClient()).thenReturn(mockClient);
@@ -167,6 +190,18 @@ void testHandleExpiredLinks()
     when(mockChunk.getStatus()).thenReturn(ChunkStatus.PENDING);
     when(mockChunkMap.get(chunkIndex)).thenReturn(mockChunk);
 
+    // Mock chunks and empty results for subsequent indices so batch continuation
+    // doesn't fail when looking up row offsets or fetching the next batch.
+    // Use lenient() since these may or may not be reached depending on timing.
+    for (long i = 2; i < TOTAL_CHUNKS; i++) {
+      ArrowResultChunk subsequentChunk = mock(ArrowResultChunk.class);
+      lenient().when(subsequentChunk.getStartRowOffset()).thenReturn(i * 100L);
+      lenient().when(mockChunkMap.get(i)).thenReturn(subsequentChunk);
+    }
+    lenient()
+        .when(mockClient.getResultChunks(eq(mockStatementId), eq(2L), anyLong()))
+        .thenReturn(buildChunkLinkFetchResult(Collections.emptyList()));
+
     ChunkLinkDownloadService<ArrowResultChunk> service =
         new ChunkLinkDownloadService<>(
             mockSession, mockStatementId, TOTAL_CHUNKS, mockChunkMap, NEXT_BATCH_START_INDEX);

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

@@ -1103,4 +1103,46 @@ public void testExecuteStatementWithClosedStatusAndNoParentStatement() throws Ex
                 null,
                 null));
   }
+
+  @Test
+  public void testGetResultChunks_DatabricksError_throwsSQLException() throws Exception {
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksSdkClient databricksSdkClient =
+        new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient);
+
+    // Simulate a 404 from the server (result expired)
+    when(apiClient.execute(any(Request.class), eq(ResultData.class)))
+        .thenThrow(new DatabricksError("404", "Results have expired", 404));
+
+    DatabricksSQLException exception =
+        assertThrows(
+            DatabricksSQLException.class,
+            () -> databricksSdkClient.getResultChunks(STATEMENT_ID, 0, 0));
+
+    assertTrue(exception.getMessage().contains("HTTP 404"));
+    assertTrue(exception.getMessage().contains("Results have expired"));
+    assertNotNull(exception.getCause());
+  }
+
+  @Test
+  public void testGetResultChunksData_DatabricksError_throwsSQLException() throws Exception {
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksSdkClient databricksSdkClient =
+        new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient);
+
+    // Simulate a 404 from the server (result expired)
+    when(apiClient.execute(any(Request.class), eq(ResultData.class)))
+        .thenThrow(new DatabricksError("404", "Results have expired", 404));
+
+    DatabricksSQLException exception =
+        assertThrows(
+            DatabricksSQLException.class,
+            () -> databricksSdkClient.getResultChunksData(STATEMENT_ID, 0));
+
+    assertTrue(exception.getMessage().contains("HTTP 404"));
+    assertTrue(exception.getMessage().contains("Results have expired"));
+    assertNotNull(exception.getCause());
+  }
 }