Fix: For fetching multiple chunks of Json chunked results (#1074)

# Fix: Fetch and assemble JSON_ARRAY results across multiple chunks

## Summary
This PR fixes an issue where JSON_ARRAY-formatted results that are split
across multiple chunks were not fully retrieved by the driver. It adds a
dedicated JsonChunkProvider to fetch and combine chunked JSON results,
wires it into InlineJsonResult, adds a client API to fetch chunk data,
and includes comprehensive unit tests.

## Motivation
Previously InlineJsonResult only used the initial response data and
ignored any additional chunks. For results returned in the `JSON_ARRAY`
format that span multiple chunks, the driver returned incomplete
datasets. This produced incorrect query results for callers when
responses were chunked by the backend.

## What this change does
- Adds `JsonChunkProvider` (new)
- Fetches and combines JSON chunks for a statement into a single
in-memory List<List<Object>>.


## Tests and validation
- Unit tests: run `JsonChunkProviderTest` and the full test suite.
- Integration/manual validation:

Author: jprakash-db

.github/workflows/prIntegrationTests.yml

@@ -16,7 +16,7 @@ jobs:
         include:
           - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!M2MAuthIntegrationTests
             fake-service-type: 'SQL_EXEC'
-          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests
+          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests,!SqlExecApiIntegrationTests
             fake-service-type: 'THRIFT_SERVER'
     steps:
       - name: Checkout PR

.github/workflows/runIntegrationTests.yml

@@ -17,7 +17,7 @@ jobs:
           - test-command: mvn -B compile test -Dtest=*IntegrationTests
             token-secret: DATABRICKS_TOKEN
             fake-service-type: 'SQL_EXEC'
-          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests
+          - test-command: mvn -B compile test -Dtest=*IntegrationTests,!M2MPrivateKeyCredentialsIntegrationTests,!SqlExecApiHybridResultsIntegrationTests,!DBFSVolumeIntegrationTests,!M2MAuthIntegrationTests,!UCVolumeIntegrationTests,!SqlExecApiIntegrationTests
             token-secret: THRIFT_DATABRICKS_TOKEN
             fake-service-type: 'THRIFT_SERVER'
     steps:

NEXT_CHANGELOG.md

@@ -17,5 +17,6 @@
 - Fix: driver failing to parse complex data types with nullable attributes.
 - Fixed: Resolved SDK token-caching regression causing token refresh on every call. SDK is now configured once to avoid excessive token endpoint hits and rate limiting.
 - Fixed: TimestampConverter.toString() returning ISO8601 format with timezone conversion instead of SQL standard format.
+- Fixed: Driver not loading complete JSON result in the case of SEA Inline without Arrow
 ---
 *Note: When making changes, please add your change under the appropriate section with a brief description.* 

src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java

@@ -57,7 +57,7 @@ private static IExecutionResult getResultHandler(
         return new ArrowStreamResult(manifest, data, statementId, session);
       case JSON_ARRAY:
         // This is used for metadata and update commands
-        return new InlineJsonResult(manifest, data);
+        return new InlineJsonResult(manifest, data, statementId, session);
       default:
         String errorMessage = String.format("Invalid response format %s", manifest.getFormat());
         LOGGER.error(errorMessage);

src/main/java/com/databricks/jdbc/api/impl/InlineJsonResult.java

@@ -1,5 +1,7 @@
 package com.databricks.jdbc.api.impl;
 
+import com.databricks.jdbc.api.internal.IDatabricksSession;
+import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.model.core.ResultManifest;
@@ -11,11 +13,22 @@ public class InlineJsonResult implements IExecutionResult {
 
   private long currentRow;
   private List<List<Object>> data;
-
+  private JsonChunkProvider chunkProvider;
   private boolean isClosed;
 
-  public InlineJsonResult(ResultManifest resultManifest, ResultData resultData) {
-    this(getDataList(resultData.getDataArray()));
+  public InlineJsonResult(
+      ResultManifest resultManifest,
+      ResultData resultData,
+      StatementId statementId,
+      IDatabricksSession session)
+      throws DatabricksSQLException {
+
+    this.chunkProvider = new JsonChunkProvider(resultManifest, resultData, statementId, session);
+    // Fetching data all at once as the data is at most 26Mb in total (SEA)
+    this.data = chunkProvider.getAllData();
+
+    this.currentRow = -1;
+    this.isClosed = false;
   }
 
   public InlineJsonResult(Object[][] rows) {
@@ -31,21 +44,6 @@ public InlineJsonResult(List<List<Object>> rows) {
     this.isClosed = false;
   }
 
-  private static List<List<Object>> getDataList(Collection<Collection<String>> dataArray) {
-    if (dataArray == null) {
-      return new ArrayList<>();
-    }
-    List<List<Object>> dataList = new ArrayList<>();
-    for (Collection<String> innerCollection : dataArray) {
-      if (innerCollection == null) {
-        dataList.add(Collections.emptyList());
-      } else {
-        dataList.add(new ArrayList<>(innerCollection));
-      }
-    }
-    return dataList;
-  }
-
   @Override
   public Object getObject(int columnIndex) throws DatabricksSQLException {
     if (isClosed()) {
@@ -86,6 +84,9 @@ public boolean hasNext() {
   public void close() {
     this.isClosed = true;
     this.data = null;
+    if (chunkProvider != null) {
+      chunkProvider.close();
+    }
   }
 
   @Override
@@ -95,7 +96,7 @@ public long getRowCount() {
 
   @Override
   public long getChunkCount() {
-    return 0;
+    return chunkProvider.getChunkCount();
   }
 
   private boolean isClosed() {

src/main/java/com/databricks/jdbc/api/impl/JsonChunkProvider.java

@@ -0,0 +1,161 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.api.internal.IDatabricksSession;
+import com.databricks.jdbc.dbclient.impl.common.StatementId;
+import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.jdbc.model.core.ResultData;
+import com.databricks.jdbc.model.core.ResultManifest;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Chunk provider for JSON_ARRAY format results that handles multiple chunks of data.
+ *
+ * <p>This provider fetches and combines data from multiple JSON chunks when the result is split
+ * across multiple chunks. It follows the same pattern as Arrow chunk providers but works with JSON
+ * data arrays instead of Arrow streams.
+ */
+public class JsonChunkProvider {
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(JsonChunkProvider.class);
+
+  private final StatementId statementId;
+  private final IDatabricksSession session;
+  private final ResultManifest resultManifest;
+  private final List<List<Object>> allData;
+  private boolean isClosed = false;
+
+  /**
+   * Creates a new JsonChunkProvider.
+   *
+   * @param resultManifest the result manifest containing chunk information
+   * @param initialResultData the first chunk of result data
+   * @param statementId the statement ID for fetching additional chunks
+   * @param session the session for making API calls
+   * @throws DatabricksSQLException if there's an error processing the initial data or fetching
+   *     chunks
+   */
+  public JsonChunkProvider(
+      ResultManifest resultManifest,
+      ResultData initialResultData,
+      StatementId statementId,
+      IDatabricksSession session)
+      throws DatabricksSQLException {
+    this.resultManifest = resultManifest;
+    this.statementId = statementId;
+    this.session = session;
+    this.allData = new ArrayList<>();
+
+    // Process all chunks
+    fetchAndCombineAllChunks(initialResultData);
+  }
+
+  /**
+   * Fetches and combines data from all chunks.
+   *
+   * @param initialResultData the first chunk of data
+   * @throws DatabricksSQLException if there's an error fetching additional chunks
+   */
+  private void fetchAndCombineAllChunks(ResultData initialResultData)
+      throws DatabricksSQLException {
+    LOGGER.debug("Starting to fetch and combine JSON chunks for statement: {}", statementId);
+
+    // Add data from the first chunk
+    addDataFromChunk(initialResultData);
+
+    // Check if there are more chunks to fetch
+    Long totalChunkCount = Objects.requireNonNullElse(resultManifest.getTotalChunkCount(), 0L);
+    if (totalChunkCount > 1) {
+      LOGGER.debug("Total chunks to fetch: {}", totalChunkCount);
+
+      // Fetch remaining chunks (starting from chunk index 1, since we already have chunk 0)
+      for (long chunkIndex = 1; chunkIndex < totalChunkCount; chunkIndex++) {
+        if (isClosed) {
+          LOGGER.warn("Provider closed while fetching chunk {}", chunkIndex);
+          break;
+        }
+
+        try {
+          LOGGER.debug("Fetching chunk {} of {}", chunkIndex, totalChunkCount);
+          ResultData chunkData =
+              session.getDatabricksClient().getResultChunksData(statementId, chunkIndex);
+          addDataFromChunk(chunkData);
+        } catch (DatabricksSQLException e) {
+          LOGGER.error("Failed to fetch chunk {} for statement {}", chunkIndex, statementId, e);
+          throw new DatabricksSQLException(
+              String.format("Failed to fetch chunk %d: %s", chunkIndex, e.getMessage()),
+              e,
+              DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR);
+        }
+      }
+    }
+
+    LOGGER.debug(
+        "Successfully combined {} total rows from {} chunks", allData.size(), totalChunkCount);
+  }
+
+  /**
+   * Adds data from a single chunk to the combined data list.
+   *
+   * @param chunkData the chunk data to add
+   */
+  private void addDataFromChunk(ResultData chunkData) {
+    if (chunkData == null || chunkData.getDataArray() == null) {
+      LOGGER.debug("Chunk data or data array is null, skipping");
+      return;
+    }
+
+    Collection<Collection<String>> dataArray = chunkData.getDataArray();
+    LOGGER.debug(
+        "Adding {} rows from chunk {}",
+        dataArray.size(),
+        chunkData.getChunkIndex() != null ? chunkData.getChunkIndex() : "unknown");
+
+    for (Collection<String> row : dataArray) {
+      if (row != null) {
+        allData.add(new ArrayList<>(row));
+      } else {
+        allData.add(new ArrayList<>());
+      }
+    }
+  }
+
+  /**
+   * Gets the combined data from all chunks.
+   *
+   * @return the list of all rows from all chunks
+   */
+  public List<List<Object>> getAllData() {
+    return allData;
+  }
+
+  /**
+   * Gets the total number of chunks that were processed.
+   *
+   * @return the chunk count
+   */
+  public long getChunkCount() {
+    LOGGER.debug("Getting total chunk count");
+    return Objects.requireNonNullElse(resultManifest.getTotalChunkCount(), 0L);
+  }
+
+  /** Closes the provider and releases resources. */
+  public void close() {
+    isClosed = true;
+    allData.clear();
+    LOGGER.debug("JsonChunkProvider closed for statement: {}", statementId);
+  }
+
+  /**
+   * Checks if the provider is closed.
+   *
+   * @return true if closed, false otherwise
+   */
+  public boolean isClosed() {
+    return isClosed;
+  }
+}

src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java

@@ -10,6 +10,7 @@
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimed;
 import com.databricks.sdk.core.DatabricksConfig;
 import java.sql.SQLException;
@@ -122,6 +123,16 @@ DatabricksResultSet getStatementResult(
   Collection<ExternalLink> getResultChunks(StatementId statementId, long chunkIndex)
       throws DatabricksSQLException;
 
+  /**
+   * Fetches the result data for given chunk index and statement-Id.
+   *
+   * @param statementId statement-Id for which chunk data should be fetched
+   * @param chunkIndex chunkIndex for which chunk data should be fetched
+   * @return ResultData containing the chunk's data array and metadata
+   */
+  ResultData getResultChunksData(StatementId statementId, long chunkIndex)
+      throws DatabricksSQLException;
+
   IDatabricksConnectionContext getConnectionContext();
 
   /**

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

@@ -432,6 +432,29 @@ public Collection<ExternalLink> getResultChunks(StatementId typedStatementId, lo
     }
   }
 
+  @Override
+  public ResultData getResultChunksData(StatementId typedStatementId, long chunkIndex)
+      throws DatabricksSQLException {
+    DatabricksThreadContextHolder.setStatementId(typedStatementId);
+    String statementId = typedStatementId.toSQLExecStatementId();
+    LOGGER.debug(
+        "public ResultData getResultChunksData(String statementId = {}, long chunkIndex = {})",
+        statementId,
+        chunkIndex);
+    GetStatementResultChunkNRequest request =
+        new GetStatementResultChunkNRequest().setStatementId(statementId).setChunkIndex(chunkIndex);
+    String path = String.format(RESULT_CHUNK_PATH, statementId, chunkIndex);
+    try {
+      Request req = new Request(Request.GET, path, apiClient.serialize(request));
+      req.withHeaders(getHeaders("getStatementResultN"));
+      return apiClient.execute(req, ResultData.class);
+    } catch (IOException e) {
+      String errorMessage = "Error while processing the get result chunk request";
+      LOGGER.error(errorMessage, e);
+      throw new DatabricksSQLException(errorMessage, e, DatabricksDriverErrorCode.SDK_CLIENT_ERROR);
+    }
+  }
+
   @Override
   public synchronized void resetAccessToken(String newAccessToken) {
     this.clientConfigurator.resetAccessTokenInConfig(newAccessToken);

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

@@ -29,6 +29,7 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.client.thrift.generated.*;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.sdk.core.DatabricksConfig;
 import com.google.common.annotations.VisibleForTesting;
@@ -327,6 +328,14 @@ public Collection<ExternalLink> getResultChunks(StatementId statementId, long ch
     return externalLinks;
   }
 
+  @Override
+  public ResultData getResultChunksData(StatementId statementId, long chunkIndex)
+      throws DatabricksSQLException {
+    throw new DatabricksSQLException(
+        "getResultChunksData method is not yet implemented for thrift client",
+        DatabricksDriverErrorCode.INVALID_STATE);
+  }
+
   @Override
   public DatabricksResultSet listTypeInfo(IDatabricksSession session) {
     LOGGER.debug("public ResultSet getTypeInfo()");