Add multi-chunk execution test (#181)

* Add multi-chunk execution test

* Store the large chunk body as separate binary file

* Improve the assertions in multi-chunk test case

Author: jayantsing-db

src/main/java/com/databricks/jdbc/core/DatabricksResultSet.java

@@ -100,7 +100,9 @@ public DatabricksResultSet(
     this.executionResult =
         ExecutionResultFactory.getResultSet(resultData, resultManifest, statementId, session);
     int rowSize = getRowCount(resultData);
-    this.resultSetMetaData = new DatabricksResultSetMetaData(statementId, resultManifest, rowSize);
+    this.resultSetMetaData =
+        new DatabricksResultSetMetaData(
+            statementId, resultManifest, rowSize, resultData.getResultLinksSize());
     this.statementType = statementType;
     this.updateCount = null;
     this.parentStatement = parentStatement;

src/main/java/com/databricks/jdbc/core/DatabricksResultSetMetaData.java

@@ -30,6 +30,7 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
   private final ImmutableList<ImmutableDatabricksColumn> columns;
   private final ImmutableMap<String, Integer> columnNameIndex;
   private final long totalRows;
+  private Long chunkCount;
   private static final String DEFAULT_CATALOGUE_NAME = "Spark";
   private static final String NULL_STRING = "null";
 
@@ -80,10 +81,11 @@ public DatabricksResultSetMetaData(
     this.columns = columnsBuilder.build();
     this.columnNameIndex = ImmutableMap.copyOf(columnNameToIndexMap);
     this.totalRows = resultManifest.getTotalRowCount();
+    this.chunkCount = resultManifest.getTotalChunkCount();
   }
 
   public DatabricksResultSetMetaData(
-      String statementId, TGetResultSetMetadataResp resultManifest, int rows) {
+      String statementId, TGetResultSetMetadataResp resultManifest, int rows, long chunkCount) {
     this.statementId = statementId;
     Map<String, Integer> columnNameToIndexMap = new HashMap<>();
     ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
@@ -108,6 +110,7 @@ public DatabricksResultSetMetaData(
     this.columns = columnsBuilder.build();
     this.columnNameIndex = ImmutableMap.copyOf(columnNameToIndexMap);
     this.totalRows = rows;
+    this.chunkCount = chunkCount;
   }
 
   public DatabricksResultSetMetaData(
@@ -283,6 +286,10 @@ public long getTotalRows() {
     return totalRows;
   }
 
+  public Long getChunkCount() {
+    return chunkCount;
+  }
+
   private ImmutableDatabricksColumn.Builder getColumnBuilder() {
     return ImmutableDatabricksColumn.builder()
         .isAutoIncrement(false)

src/test/java/com/databricks/jdbc/client/impl/thrift/commons/DatabricksThriftAccessorTest.java

@@ -18,6 +18,7 @@
 import com.databricks.sdk.core.DatabricksConfig;
 import com.databricks.sdk.service.sql.StatementState;
 import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.Collections;
 import org.apache.thrift.TException;
 import org.apache.thrift.protocol.TProtocol;
@@ -46,11 +47,16 @@ public class DatabricksThriftAccessorTest {
           .setMaxBytes(DEFAULT_BYTE_LIMIT);
   private static final TGetResultSetMetadataReq resultSetMetadataReq =
       new TGetResultSetMetadataReq().setOperationHandle(tOperationHandle);
+
+  private static final TRowSet rowSet = new TRowSet().setResultLinks(new ArrayList<>(2));
+
   private static final TFetchResultsResp response =
       new TFetchResultsResp()
           .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))
           .setResultSetMetadata(
-              new TGetResultSetMetadataResp().setResultFormat(TSparkRowSetType.COLUMN_BASED_SET));
+              new TGetResultSetMetadataResp().setResultFormat(TSparkRowSetType.COLUMN_BASED_SET))
+          .setResults(rowSet);
+
   private static final TSparkDirectResults directResults =
       new TSparkDirectResults()
           .setResultSet(response)

src/test/java/com/databricks/jdbc/core/DatabricksResultSetMetaDataTest.java

@@ -101,7 +101,7 @@ public void testColumnsForVolumeOperation() throws SQLException {
   @Test
   public void testThriftColumns() throws SQLException {
     DatabricksResultSetMetaData metaData =
-        new DatabricksResultSetMetaData(STATEMENT_ID, getThriftResultManifest(), 10);
+        new DatabricksResultSetMetaData(STATEMENT_ID, getThriftResultManifest(), 10, 1);
     Assertions.assertEquals(10, metaData.getTotalRows());
     Assertions.assertEquals(1, metaData.getColumnCount());
     Assertions.assertEquals("testCol", metaData.getColumnName(1));
@@ -111,7 +111,7 @@ public void testThriftColumns() throws SQLException {
   public void testEmptyAndNullThriftColumns() throws SQLException {
     TGetResultSetMetadataResp resultSetMetadataResp = new TGetResultSetMetadataResp();
     DatabricksResultSetMetaData metaData =
-        new DatabricksResultSetMetaData(STATEMENT_ID, resultSetMetadataResp, 0);
+        new DatabricksResultSetMetaData(STATEMENT_ID, resultSetMetadataResp, 0, 1);
     Assertions.assertEquals(0, metaData.getColumnCount());
 
     resultSetMetadataResp.setSchema(new TTableSchema());

src/test/java/com/databricks/jdbc/integration/fakeservice/FakeServiceExtension.java

@@ -2,20 +2,30 @@
 
 import static com.databricks.jdbc.driver.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX;
 import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
+import static com.github.tomakehurst.wiremock.common.AbstractFileSource.byFileExtension;
 
 import com.databricks.jdbc.client.http.DatabricksHttpClient;
 import com.databricks.jdbc.driver.DatabricksJdbcConstants.FakeServiceType;
 import com.databricks.jdbc.integration.IntegrationTestUtil;
+import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
+import com.github.tomakehurst.wiremock.common.BinaryFile;
+import com.github.tomakehurst.wiremock.common.ContentTypes;
 import com.github.tomakehurst.wiremock.common.SingleRootFileSource;
+import com.github.tomakehurst.wiremock.common.TextFile;
+import com.github.tomakehurst.wiremock.http.ContentTypeHeader;
+import com.github.tomakehurst.wiremock.http.HttpHeaders;
 import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo;
 import com.github.tomakehurst.wiremock.recording.RecordSpecBuilder;
 import com.github.tomakehurst.wiremock.standalone.JsonFileMappingsSource;
 import com.github.tomakehurst.wiremock.stubbing.StubMapping;
+import com.github.tomakehurst.wiremock.stubbing.StubMappingCollection;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -67,6 +77,9 @@ public class FakeServiceExtension extends DatabricksWireMockExtension {
    */
   private static final long MAX_STUBBING_BINARY_SIZE = 102400;
 
+  /** Root directory for extracted body files. */
+  private static final String EXTRACTED_BODY_FILE_ROOT = "src/test/resources/__files";
+
   /**
    * Environment variable holding the fake service mode.
    *
@@ -195,9 +208,24 @@ private String getStubbingDir(ExtensionContext context) {
   }
 
   /** Loads stub mappings from the stubbing directory. */
-  private void loadStubMappings(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionContext context) {
-    String stubbingDir = getStubbingDir(context);
-    wireMockRuntimeInfo.getWireMock().loadMappingsFrom(stubbingDir);
+  private void loadStubMappings(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionContext context)
+      throws IOException {
+    final String stubbingDir = getStubbingDir(context);
+    final SingleRootFileSource fileSource = new SingleRootFileSource(stubbingDir + "/mappings");
+
+    final List<TextFile> mappingFiles =
+        fileSource.listFilesRecursively().stream()
+            .filter(byFileExtension("json"))
+            .collect(Collectors.toList());
+
+    for (TextFile mappingFile : mappingFiles) {
+      final StubMappingCollection stubCollection =
+          JsonUtils.read(mappingFile.readContents(), StubMappingCollection.class);
+      for (StubMapping mapping : stubCollection.getMappingOrMappings()) {
+        embedExtractedBodyFile(mapping);
+        wireMockRuntimeInfo.getWireMock().register(mapping);
+      }
+    }
   }
 
   /** Starts recording stub mappings. */
@@ -241,6 +269,34 @@ private void clearFakeServiceProperties() {
     System.clearProperty(targetBaseUrl + FAKE_SERVICE_URI_PROP_SUFFIX);
   }
 
+  /** Embeds the extracted body file content into the stub mapping. */
+  private static void embedExtractedBodyFile(final StubMapping mapping) {
+    final SingleRootFileSource fileSource = new SingleRootFileSource(EXTRACTED_BODY_FILE_ROOT);
+    final String bodyFileName = mapping.getResponse().getBodyFileName();
+
+    if (bodyFileName != null) {
+      final ResponseDefinitionBuilder responseDefinitionBuilder =
+          ResponseDefinitionBuilder.like(mapping.getResponse()).withBodyFile(null);
+      if (ContentTypes.determineIsTextFromMimeType(getMimeType(mapping))) {
+        final TextFile bodyFile = fileSource.getTextFileNamed(bodyFileName);
+        responseDefinitionBuilder.withBody(bodyFile.readContentsAsString());
+      } else {
+        BinaryFile bodyFile = fileSource.getBinaryFileNamed(bodyFileName);
+        responseDefinitionBuilder.withBody(bodyFile.readContents());
+      }
+
+      mapping.setResponse(responseDefinitionBuilder.build());
+    }
+  }
+
+  /** Gets the MIME type of the response body. */
+  private static String getMimeType(StubMapping mapping) {
+    return Optional.ofNullable(mapping.getResponse().getHeaders())
+        .map(HttpHeaders::getContentTypeHeader)
+        .map(ContentTypeHeader::mimeTypePart)
+        .orElse(null);
+  }
+
   /** Deletes files in the given directory. */
   private static void deleteFilesInDir(String dirPath) throws IOException {
     Path dir = Paths.get(dirPath);

src/test/java/com/databricks/jdbc/integration/fakeservice/JsonUtils.java

@@ -0,0 +1,61 @@
+package com.databricks.jdbc.integration.fakeservice;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.core.*;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.ObjectWriter;
+import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.databind.cfg.JsonNodeFeature;
+import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
+import java.io.IOException;
+
+/** Utility class for JSON operations. */
+public class JsonUtils {
+
+  /** {@link ObjectMapper} instance holder. */
+  private static final InheritableThreadLocal<ObjectMapper> objectMapperHolder =
+      new InheritableThreadLocal<>() {
+        @Override
+        protected ObjectMapper initialValue() {
+          ObjectMapper objectMapper =
+              new ObjectMapper(
+                  new JsonFactoryBuilder()
+                      .streamReadConstraints(
+                          StreamReadConstraints.builder()
+                              .maxStringLength(Math.toIntExact(30000000))
+                              .build())
+                      .build());
+          objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
+          objectMapper.configure(JsonNodeFeature.STRIP_TRAILING_BIGDECIMAL_ZEROES, false);
+          objectMapper.configure(JsonParser.Feature.ALLOW_COMMENTS, true);
+          objectMapper.configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true);
+          objectMapper.configure(JsonParser.Feature.IGNORE_UNDEFINED, true);
+          objectMapper.configure(JsonGenerator.Feature.WRITE_BIGDECIMAL_AS_PLAIN, true);
+          objectMapper.configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true);
+          objectMapper.registerModule(new JavaTimeModule());
+          objectMapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
+          objectMapper.enable(JsonParser.Feature.INCLUDE_SOURCE_IN_LOCATION);
+          return objectMapper;
+        }
+      };
+
+  /** Read the given JSON string and return the object of the given class. */
+  public static <T> T read(String json, Class<T> clazz) throws JsonProcessingException {
+    return objectMapperHolder.get().readValue(json, clazz);
+  }
+
+  /** Read the given byte array and return the object of the given class. */
+  public static <T> T read(byte[] stream, Class<T> clazz) throws IOException {
+    return objectMapperHolder.get().readValue(stream, clazz);
+  }
+
+  /** Write the given object to a JSON string. */
+  public static <T> String write(T object, Class<?> view) throws JsonProcessingException {
+    ObjectWriter objectWriter = objectMapperHolder.get().writerWithDefaultPrettyPrinter();
+    if (view != null) {
+      objectWriter = objectWriter.withView(view);
+    }
+    return objectWriter.writeValueAsString(object);
+  }
+}

src/test/java/com/databricks/jdbc/integration/fakeservice/StubMappingCredentialsCleaner.java

@@ -2,7 +2,9 @@
 
 import static com.databricks.jdbc.client.impl.sdk.PathConstants.STATEMENT_PATH;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.github.tomakehurst.wiremock.common.FileSource;
+import com.github.tomakehurst.wiremock.common.Json;
 import com.github.tomakehurst.wiremock.extension.Parameters;
 import com.github.tomakehurst.wiremock.extension.StubMappingTransformer;
 import com.github.tomakehurst.wiremock.stubbing.StubMapping;
@@ -34,13 +36,17 @@ public StubMapping transform(StubMapping stubMapping, FileSource files, Paramete
     String serverHeaderValue =
         stubMapping.getResponse().getHeaders().getHeader(SERVER_HEADER_NAME).firstValue();
 
-    if (STATEMENT_PATH.equals(requestUrl)
-        || AMAZON_S3_SERVER_VALUE.equals(serverHeaderValue)
+    if (AMAZON_S3_SERVER_VALUE.equals(serverHeaderValue)
+        || requestUrl.startsWith(STATEMENT_PATH)
         || serverHeaderValue.startsWith(AZURE_STORAGE_SERVER_VALUE)) {
       // Clean credentials from statement requests (embedded S3 links) and Amazon S3 responses.
-      String jsonString = StubMapping.buildJsonStringFor(stubMapping);
-      String transformedJsonString = removeSensitiveCredentials(jsonString);
-      return StubMapping.buildFrom(transformedJsonString);
+      try {
+        final String jsonString = getJsonStringFromStubMapping(stubMapping);
+        final String transformedJsonString = removeSensitiveCredentials(jsonString);
+        return getStubMappingFromJsonString(transformedJsonString);
+      } catch (JsonProcessingException e) {
+        throw new RuntimeException(e);
+      }
     } else {
       return stubMapping;
     }
@@ -64,4 +70,16 @@ private String removeSensitiveCredentials(String jsonString) {
 
     return buffer.toString();
   }
+
+  /** Returns the JSON string representation of the given {@link StubMapping}. */
+  private String getJsonStringFromStubMapping(StubMapping stubMapping)
+      throws JsonProcessingException {
+    return JsonUtils.write(stubMapping, Json.PublicView.class);
+  }
+
+  /** Returns the {@link StubMapping} from the given JSON string. */
+  private StubMapping getStubMappingFromJsonString(String jsonString)
+      throws JsonProcessingException {
+    return JsonUtils.read(jsonString, StubMapping.class);
+  }
 }