Adding GET extended command for UC volume operation (#384)

Author: gopalldb

src/main/java/com/databricks/jdbc/client/IDatabricksUCVolumeClient.java

@@ -81,6 +81,18 @@ boolean getObject(
       String catalog, String schema, String volume, String objectPath, String localPath)
       throws SQLException;
 
+  /**
+   * getObject(): Retrieves an object as input stream from the UC Volume
+   *
+   * @param catalog the catalog name of the cloud storage
+   * @param schema the schema name of the cloud storage
+   * @param volume the UC volume name of the cloud storage
+   * @param objectPath the path of the object (file) from the volume as the root directory
+   * @return an instance of input stream
+   */
+  InputStream getObject(String catalog, String schema, String volume, String objectPath)
+      throws SQLException;
+
   /**
    * putObject(): Upload data from a local path to a specified path within a UC Volume.
    *

src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksUCVolumeClient.java

@@ -6,7 +6,8 @@
 import com.databricks.jdbc.client.IDatabricksUCVolumeClient;
 import com.databricks.jdbc.commons.LogLevel;
 import com.databricks.jdbc.commons.util.LoggingUtil;
-import com.databricks.jdbc.core.DatabricksStatement;
+import com.databricks.jdbc.core.IDatabricksResultSet;
+import com.databricks.jdbc.core.IDatabricksStatement;
 import java.io.InputStream;
 import java.sql.*;
 import java.util.ArrayList;
@@ -41,6 +42,12 @@ private String createGetObjectQuery(
         "GET '/Volumes/%s/%s/%s/%s' TO '%s'", catalog, schema, volume, objectPath, localPath);
   }
 
+  private String createGetObjectQueryForInputStream(
+      String catalog, String schema, String volume, String objectPath) {
+    return String.format(
+        "GET '/Volumes/%s/%s/%s/%s' TO '__input_stream__'", catalog, schema, volume, objectPath);
+  }
+
   private String createPutObjectQuery(
       String catalog,
       String schema,
@@ -293,6 +300,36 @@ public boolean getObject(
     return volumeOperationStatus;
   }
 
+  @Override
+  public InputStream getObject(String catalog, String schema, String volume, String objectPath)
+      throws SQLException {
+
+    LoggingUtil.log(
+        LogLevel.DEBUG,
+        String.format(
+            "Entering getObject method with parameters: catalog={%s}, schema={%s}, volume={%s}, objectPath={%s}",
+            catalog, schema, volume, objectPath));
+
+    String getObjectQuery = createGetObjectQueryForInputStream(catalog, schema, volume, objectPath);
+
+    Statement statement = connection.createStatement();
+    IDatabricksStatement databricksStatement = (IDatabricksStatement) statement;
+    databricksStatement.allowInputStreamForVolumeOperation(true);
+
+    ResultSet resultSet = statement.executeQuery(getObjectQuery);
+    LoggingUtil.log(LogLevel.INFO, "GET query executed successfully");
+
+    try {
+      if (resultSet.next()) {
+        return ((IDatabricksResultSet) resultSet).getVolumeOperationInputStream();
+      }
+      return null;
+    } catch (SQLException e) {
+      LoggingUtil.log(LogLevel.ERROR, "GET query execution failed " + e);
+      throw e;
+    }
+  }
+
   public boolean putObject(
       String catalog,
       String schema,
@@ -353,11 +390,11 @@ public boolean putObject(
     boolean volumeOperationStatus = false;
 
     try (Statement statement = connection.createStatement()) {
-      DatabricksStatement databricksStatement = (DatabricksStatement) statement;
+      IDatabricksStatement databricksStatement = (IDatabricksStatement) statement;
       databricksStatement.allowInputStreamForVolumeOperation(true);
       databricksStatement.setInputStreamForUCVolume(inputStream);
 
-      ResultSet resultSet = databricksStatement.executeQuery(putObjectQueryForInputStream);
+      ResultSet resultSet = statement.executeQuery(putObjectQueryForInputStream);
       LoggingUtil.log(LogLevel.INFO, "PUT query executed successfully");
 
       if (resultSet.next()) {

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

@@ -30,6 +30,7 @@ public class DatabricksUCVolumeClientTest {
   @Mock Statement statement;
 
   @Mock DatabricksStatement databricksStatement;
+  @Mock DatabricksResultSet databricksResultSet;
 
   @Mock ResultSet resultSet;
   @Mock ResultSet resultSet_abc_volume1;
@@ -523,6 +524,37 @@ private static Stream<Arguments> provideParametersForDeleteObject_InvalidObjectP
         Arguments.of("test_catalog", "test_schema", "test_volume", "invalid_objectpath"));
   }
 
+  @ParameterizedTest
+  @MethodSource("provideParametersForGetObjectWithInputStream")
+  public void testGetObjectWithInputStream(
+      String catalog, String schema, String volume, String objectPath, InputStream expected)
+      throws SQLException {
+    DatabricksUCVolumeClient client = new DatabricksUCVolumeClient(connection);
+
+    when(connection.createStatement()).thenReturn(databricksStatement);
+
+    String getObjectQuery =
+        String.format(
+            "GET '/Volumes/%s/%s/%s/%s' TO '__input_stream__'",
+            catalog, schema, volume, objectPath);
+    when(databricksStatement.executeQuery(getObjectQuery)).thenReturn(databricksResultSet);
+    when(databricksResultSet.next()).thenReturn(true);
+    when(databricksResultSet.getVolumeOperationInputStream()).thenReturn(expected);
+
+    InputStream result = client.getObject(catalog, schema, volume, objectPath);
+
+    assertEquals(expected, result);
+    verify(databricksStatement).executeQuery(getObjectQuery);
+    verify(databricksStatement).allowInputStreamForVolumeOperation(true);
+  }
+
+  private static Stream<Arguments> provideParametersForGetObjectWithInputStream() {
+    InputStream inputStream =
+        new ByteArrayInputStream("test data".getBytes(StandardCharsets.UTF_8));
+    return Stream.of(
+        Arguments.of("test_catalog", "test_schema", "test_volume", "test_objectpath", inputStream));
+  }
+
   @ParameterizedTest
   @MethodSource("provideParametersForPutObjectWithInputStream")
   public void testPutObjectWithInputStream(
@@ -551,6 +583,7 @@ public void testPutObjectWithInputStream(
 
     assertEquals(expected, result);
     verify(databricksStatement).executeQuery(putObjectQuery);
+    verify(databricksStatement).allowInputStreamForVolumeOperation(true);
   }
 
   private static Stream<Arguments> provideParametersForPutObjectWithInputStream() {