Skip to content

Commit c19eaa6

Browse files
committed
Fix UC vol commands
1 parent 078f53f commit c19eaa6

File tree

4 files changed

+25
-8
lines changed

4 files changed

+25
-8
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ public static Object convert(Object object, ColumnInfoTypeName requiredType)
7171
// Struct and Array are present in Arrow data in the VARCHAR ValueVector format
7272
case STRUCT:
7373
case ARRAY:
74+
case MAP:
7475
return convertToString(object);
7576
case DATE:
7677
return convertToDate(object);

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -532,7 +532,10 @@ protected static boolean shouldReturnResultSet(String query) {
532532
|| UNION_PATTERN.matcher(trimmedQuery).find()
533533
|| INTERSECT_PATTERN.matcher(trimmedQuery).find()
534534
|| EXCEPT_PATTERN.matcher(trimmedQuery).find()
535-
|| DECLARE_PATTERN.matcher(trimmedQuery).find()) {
535+
|| DECLARE_PATTERN.matcher(trimmedQuery).find()
536+
|| PUT_PATTERN.matcher(trimmedQuery).find()
537+
|| GET_PATTERN.matcher(trimmedQuery).find()
538+
|| REMOVE_PATTERN.matcher(trimmedQuery).find()) {
536539
return true;
537540
}
538541

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class VolumeOperationResult implements IExecutionResult {
6060
private void initHandler(IExecutionResult resultHandler) throws DatabricksSQLException {
6161
String operation = getString(resultHandler.getObject(0));
6262
String presignedUrl = getString(resultHandler.getObject(1));
63-
String localFile = getString(resultHandler.getObject(3));
63+
String localFile = columnCount > 3 ? getString(resultHandler.getObject(3)) : null;
6464
Map<String, String> headers = getHeaders(getString(resultHandler.getObject(2)));
6565
this.volumeOperationExecutor =
6666
new VolumeOperationExecutor(
@@ -83,11 +83,18 @@ private String getString(Object obj) {
8383

8484
private Map<String, String> getHeaders(String headersVal) throws DatabricksSQLException {
8585
if (headersVal != null && !headersVal.isEmpty()) {
86-
ObjectMapper objectMapper = new ObjectMapper();
87-
try {
88-
return objectMapper.readValue(headersVal, Map.class);
89-
} catch (JsonProcessingException e) {
90-
throw new DatabricksSQLException("Failed to parse headers", e);
86+
// Map is encoded in extra [] while doing toString
87+
String headers =
88+
headersVal.charAt(0) == '['
89+
? headersVal.substring(1, headersVal.length() - 1)
90+
: headersVal;
91+
if (!headers.isEmpty()) {
92+
ObjectMapper objectMapper = new ObjectMapper();
93+
try {
94+
return objectMapper.readValue(headers, Map.class);
95+
} catch (JsonProcessingException e) {
96+
throw new DatabricksSQLException("Failed to parse headers", e);
97+
}
9198
}
9299
}
93100
return new HashMap<>();
@@ -111,7 +118,7 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
111118
if (currentRowIndex < 0) {
112119
throw new DatabricksSQLException("Invalid row access");
113120
}
114-
if (columnIndex == 1) {
121+
if (columnIndex == 0) {
115122
return volumeOperationExecutor.getStatus().name();
116123
} else {
117124
throw new DatabricksSQLException("Invalid column access");

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,12 @@ public enum FakeServiceType {
246246
Pattern.compile("\\s+EXCEPT\\s+", Pattern.CASE_INSENSITIVE);
247247
public static final Pattern DECLARE_PATTERN =
248248
Pattern.compile("^(\\s*\\()*\\s*DECLARE", Pattern.CASE_INSENSITIVE);
249+
public static final Pattern PUT_PATTERN =
250+
Pattern.compile("^(\\s*\\()*\\s*GET", Pattern.CASE_INSENSITIVE);
251+
public static final Pattern GET_PATTERN =
252+
Pattern.compile("^(\\s*\\()*\\s*PUT", Pattern.CASE_INSENSITIVE);
253+
public static final Pattern REMOVE_PATTERN =
254+
Pattern.compile("^(\\s*\\()*\\s*REMOVE", Pattern.CASE_INSENSITIVE);
249255
static final int DBSQL_MIN_MAJOR_VERSION_FOR_NEW_METADATA = 2024;
250256
static final int DBSQL_MIN_MINOR_VERSION_FOR_NEW_METADATA = 30;
251257

0 commit comments

Comments
 (0)