Merge pull request #298 from databricks/vol-5-jul

Minor Fixes in UC volume commands client side changes

Author: gopalldb

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

@@ -71,6 +71,7 @@ public static Object convert(Object object, ColumnInfoTypeName requiredType)
         // Struct and Array are present in Arrow data in the VARCHAR ValueVector format
       case STRUCT:
       case ARRAY:
+      case MAP:
         return convertToString(object);
       case DATE:
         return convertToDate(object);

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

@@ -532,7 +532,10 @@ protected static boolean shouldReturnResultSet(String query) {
         || UNION_PATTERN.matcher(trimmedQuery).find()
         || INTERSECT_PATTERN.matcher(trimmedQuery).find()
         || EXCEPT_PATTERN.matcher(trimmedQuery).find()
-        || DECLARE_PATTERN.matcher(trimmedQuery).find()) {
+        || DECLARE_PATTERN.matcher(trimmedQuery).find()
+        || PUT_PATTERN.matcher(trimmedQuery).find()
+        || GET_PATTERN.matcher(trimmedQuery).find()
+        || REMOVE_PATTERN.matcher(trimmedQuery).find()) {
       return true;
     }
 

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

@@ -60,7 +60,7 @@ class VolumeOperationResult implements IExecutionResult {
   private void initHandler(IExecutionResult resultHandler) throws DatabricksSQLException {
     String operation = getString(resultHandler.getObject(0));
     String presignedUrl = getString(resultHandler.getObject(1));
-    String localFile = getString(resultHandler.getObject(3));
+    String localFile = columnCount > 3 ? getString(resultHandler.getObject(3)) : null;
     Map<String, String> headers = getHeaders(getString(resultHandler.getObject(2)));
     this.volumeOperationExecutor =
         new VolumeOperationExecutor(
@@ -83,11 +83,18 @@ private String getString(Object obj) {
 
   private Map<String, String> getHeaders(String headersVal) throws DatabricksSQLException {
     if (headersVal != null && !headersVal.isEmpty()) {
-      ObjectMapper objectMapper = new ObjectMapper();
-      try {
-        return objectMapper.readValue(headersVal, Map.class);
-      } catch (JsonProcessingException e) {
-        throw new DatabricksSQLException("Failed to parse headers", e);
+      // Map is encoded in extra [] while doing toString
+      String headers =
+          headersVal.charAt(0) == '['
+              ? headersVal.substring(1, headersVal.length() - 1)
+              : headersVal;
+      if (!headers.isEmpty()) {
+        ObjectMapper objectMapper = new ObjectMapper();
+        try {
+          return objectMapper.readValue(headers, Map.class);
+        } catch (JsonProcessingException e) {
+          throw new DatabricksSQLException("Failed to parse headers", e);
+        }
       }
     }
     return new HashMap<>();
@@ -111,7 +118,7 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
     if (currentRowIndex < 0) {
       throw new DatabricksSQLException("Invalid row access");
     }
-    if (columnIndex == 1) {
+    if (columnIndex == 0) {
       return volumeOperationExecutor.getStatus().name();
     } else {
       throw new DatabricksSQLException("Invalid column access");

src/main/java/com/databricks/jdbc/driver/DatabricksJdbcConstants.java

@@ -247,6 +247,12 @@ public enum FakeServiceType {
       Pattern.compile("\\s+EXCEPT\\s+", Pattern.CASE_INSENSITIVE);
   public static final Pattern DECLARE_PATTERN =
       Pattern.compile("^(\\s*\\()*\\s*DECLARE", Pattern.CASE_INSENSITIVE);
+  public static final Pattern PUT_PATTERN =
+      Pattern.compile("^(\\s*\\()*\\s*GET", Pattern.CASE_INSENSITIVE);
+  public static final Pattern GET_PATTERN =
+      Pattern.compile("^(\\s*\\()*\\s*PUT", Pattern.CASE_INSENSITIVE);
+  public static final Pattern REMOVE_PATTERN =
+      Pattern.compile("^(\\s*\\()*\\s*REMOVE", Pattern.CASE_INSENSITIVE);
   static final int DBSQL_MIN_MAJOR_VERSION_FOR_NEW_METADATA = 2024;
   static final int DBSQL_MIN_MINOR_VERSION_FOR_NEW_METADATA = 30;
 

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

@@ -65,7 +65,7 @@ public void testGetResult_Get() throws Exception {
     assertEquals(-1, volumeOperationResult.getCurrentRow());
     assertTrue(volumeOperationResult.next());
     assertEquals(0, volumeOperationResult.getCurrentRow());
-    assertEquals("SUCCEEDED", volumeOperationResult.getObject(1));
+    assertEquals("SUCCEEDED", volumeOperationResult.getObject(0));
     assertFalse(volumeOperationResult.hasNext());
     assertFalse(volumeOperationResult.next());
 
@@ -263,7 +263,7 @@ public void testGetResult_Put() throws Exception {
     assertEquals(-1, volumeOperationResult.getCurrentRow());
     assertTrue(volumeOperationResult.next());
     assertEquals(0, volumeOperationResult.getCurrentRow());
-    assertEquals("SUCCEEDED", volumeOperationResult.getObject(1));
+    assertEquals("SUCCEEDED", volumeOperationResult.getObject(0));
     assertFalse(volumeOperationResult.hasNext());
     assertFalse(volumeOperationResult.next());
     assertTrue(file.delete());
@@ -386,7 +386,7 @@ public void testGetResult_Remove() throws Exception {
     assertEquals(-1, volumeOperationResult.getCurrentRow());
     assertTrue(volumeOperationResult.next());
     assertEquals(0, volumeOperationResult.getCurrentRow());
-    assertEquals("SUCCEEDED", volumeOperationResult.getObject(1));
+    assertEquals("SUCCEEDED", volumeOperationResult.getObject(0));
     assertFalse(volumeOperationResult.hasNext());
     assertFalse(volumeOperationResult.next());
     try {