[ES-1006739] Input stream support for UC Volume Operation implementation (#382)

* changes

* Minor test changes

* Updated changes

* Address comemnts

* minor fix

Author: gopalldb

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

@@ -55,7 +55,8 @@ public DatabricksResultSet(
     this.statementStatus = statementStatus;
     this.statementId = statementId;
     this.executionResult =
-        ExecutionResultFactory.getResultSet(resultData, resultManifest, statementId, session);
+        ExecutionResultFactory.getResultSet(
+            resultData, resultManifest, statementId, session, parentStatement, this);
     this.resultSetMetaData = new DatabricksResultSetMetaData(statementId, resultManifest);
     this.statementType = statementType;
     this.updateCount = null;
@@ -99,7 +100,8 @@ public DatabricksResultSet(
     }
     this.statementId = statementId;
     this.executionResult =
-        ExecutionResultFactory.getResultSet(resultData, resultManifest, statementId, session);
+        ExecutionResultFactory.getResultSet(
+            resultData, resultManifest, statementId, session, parentStatement, this);
     long rowSize = getRowCount(resultData);
     this.resultSetMetaData =
         new DatabricksResultSetMetaData(

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

@@ -1,16 +1,21 @@
 package com.databricks.jdbc.core;
 
+import com.databricks.sdk.service.sql.StatementStatus;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.Reader;
 import java.math.BigDecimal;
 import java.net.URL;
 import java.sql.*;
 import java.util.Calendar;
 import java.util.Map;
+import org.apache.http.HttpEntity;
+import org.apache.http.entity.InputStreamEntity;
 
 /** Empty implementation of ResultSet */
-class EmptyResultSet implements ResultSet {
+class EmptyResultSet implements ResultSet, IDatabricksResultSet {
   private boolean isClosed;
+  private HttpEntity httpEntity = null;
 
   EmptyResultSet() {
     isClosed = false;
@@ -1082,4 +1087,39 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
     checkIfClosed();
     return false;
   }
+
+  @Override
+  public String statementId() {
+    return null;
+  }
+
+  @Override
+  public StatementStatus getStatementStatus() {
+    return null;
+  }
+
+  @Override
+  public long getUpdateCount() throws SQLException {
+    return 0;
+  }
+
+  @Override
+  public boolean hasUpdateCount() throws SQLException {
+    return false;
+  }
+
+  @Override
+  public void setVolumeOperationEntityStream(HttpEntity httpEntity)
+      throws SQLException, IOException {
+    this.httpEntity = httpEntity;
+  }
+
+  @Override
+  public InputStreamEntity getVolumeOperationInputStream() throws SQLException {
+    try {
+      return new InputStreamEntity(httpEntity.getContent(), httpEntity.getContentLength());
+    } catch (IOException e) {
+      return null;
+    }
+  }
 }

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

@@ -11,15 +11,22 @@
 
 class ExecutionResultFactory {
   static IExecutionResult getResultSet(
-      ResultData data, ResultManifest manifest, String statementId, IDatabricksSession session) {
+      ResultData data,
+      ResultManifest manifest,
+      String statementId,
+      IDatabricksSession session,
+      IDatabricksStatement statement,
+      IDatabricksResultSet resultSet) {
     IExecutionResult resultHandler = getResultHandler(data, manifest, statementId, session);
     if (manifest.getIsVolumeOperation() != null && manifest.getIsVolumeOperation()) {
       return new VolumeOperationResult(
           statementId,
           manifest.getTotalRowCount(),
           manifest.getSchema().getColumnCount(),
           session,
-          resultHandler);
+          resultHandler,
+          statement,
+          resultSet);
     } else {
       return resultHandler;
     }
@@ -46,7 +53,9 @@ static IExecutionResult getResultSet(
       TRowSet data,
       TGetResultSetMetadataResp manifest,
       String statementId,
-      IDatabricksSession session)
+      IDatabricksSession session,
+      IDatabricksStatement statement,
+      IDatabricksResultSet resultSet)
       throws DatabricksSQLException {
     IExecutionResult resultHandler = getResultHandler(data, manifest, statementId, session);
     if (manifest.isSetIsStagingOperation() && manifest.isIsStagingOperation()) {
@@ -55,7 +64,9 @@ static IExecutionResult getResultSet(
           DatabricksThriftHelper.getRowCount(data),
           manifest.getSchema().getColumnsSize(),
           session,
-          resultHandler);
+          resultHandler,
+          statement,
+          resultSet);
     } else {
       return resultHandler;
     }

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

@@ -5,6 +5,7 @@
 import com.databricks.jdbc.commons.LogLevel;
 import com.databricks.jdbc.commons.util.LoggingUtil;
 import java.io.*;
+import java.sql.SQLException;
 import java.util.*;
 import org.apache.http.HttpEntity;
 import org.apache.http.client.methods.CloseableHttpResponse;
@@ -13,6 +14,7 @@
 import org.apache.http.client.methods.HttpPut;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.FileEntity;
+import org.apache.http.entity.InputStreamEntity;
 import org.apache.http.util.EntityUtils;
 
 /** Executor for volume operations */
@@ -30,8 +32,10 @@ class VolumeOperationExecutor implements Runnable {
   private final String localFilePath;
   private final Map<String, String> headers;
   private final Set<String> allowedVolumeIngestionPaths;
+  private final IDatabricksStatement statement;
+  private final IDatabricksResultSet resultSet;
+  private final IDatabricksHttpClient databricksHttpClient;
   private VolumeOperationStatus status;
-  private IDatabricksHttpClient databricksHttpClient;
   private String errorMessage;
 
   VolumeOperationExecutor(
@@ -40,13 +44,17 @@ class VolumeOperationExecutor implements Runnable {
       Map<String, String> headers,
       String localFilePath,
       String allowedVolumeIngestionPathString,
-      IDatabricksHttpClient databricksHttpClient) {
+      IDatabricksHttpClient databricksHttpClient,
+      IDatabricksStatement statement,
+      IDatabricksResultSet resultSet) {
     this.operationType = operationType;
     this.operationUrl = operationUrl;
     this.localFilePath = localFilePath;
     this.headers = headers;
     this.allowedVolumeIngestionPaths = getAllowedPaths(allowedVolumeIngestionPathString);
     this.databricksHttpClient = databricksHttpClient;
+    this.statement = statement;
+    this.resultSet = resultSet;
     this.status = VolumeOperationStatus.PENDING;
     this.errorMessage = null;
   }
@@ -101,6 +109,16 @@ String getErrorMessage() {
   }
 
   private void validateLocalFilePath() {
+    try {
+      if (statement.isAllowedInputStreamForVolumeOperation()) {
+        return;
+      }
+    } catch (DatabricksSQLException e) {
+      status = VolumeOperationStatus.ABORTED;
+      errorMessage = "Volume operation called on closed statement: " + e.getMessage();
+      LoggingUtil.log(LogLevel.ERROR, errorMessage);
+      return;
+    }
     if (allowedVolumeIngestionPaths.isEmpty()) {
       LoggingUtil.log(LogLevel.ERROR, "Volume ingestion paths are not set");
       status = VolumeOperationStatus.ABORTED;
@@ -136,6 +154,36 @@ private void executeGetOperation() {
     HttpGet httpGet = new HttpGet(operationUrl);
     headers.forEach(httpGet::addHeader);
 
+    HttpEntity entity = null;
+    try {
+      // We return the input stream directly to clients, if they want to consume as input stream
+      if (statement.isAllowedInputStreamForVolumeOperation()) {
+        CloseableHttpResponse response = databricksHttpClient.execute(httpGet);
+        if (!isSuccessfulHttpResponse(response)) {
+          status = VolumeOperationStatus.FAILED;
+          errorMessage =
+              String.format(
+                  "Failed to fetch content from volume with error code {%s} for input stream and error {%s}",
+                  response.getStatusLine().getStatusCode(),
+                  response.getStatusLine().getReasonPhrase());
+          LoggingUtil.log(LogLevel.ERROR, errorMessage);
+          return;
+        }
+        entity = response.getEntity();
+        if (entity != null) {
+          this.resultSet.setVolumeOperationEntityStream(entity);
+        }
+        status = VolumeOperationStatus.SUCCEEDED;
+        return;
+      }
+    } catch (SQLException | IOException e) {
+      status = VolumeOperationStatus.FAILED;
+      errorMessage = "Failed to execute GET operation for input stream: " + e.getMessage();
+      LoggingUtil.log(LogLevel.ERROR, errorMessage);
+      return;
+    }
+
+    // Copy the data in local file as requested by user
     File localFile = new File(localFilePath);
     if (localFile.exists()) {
       LoggingUtil.log(
@@ -151,13 +199,13 @@ private void executeGetOperation() {
         LoggingUtil.log(
             LogLevel.ERROR,
             String.format(
-                "Failed to fetch content from volume with error {} for local file {%s}",
+                "Failed to fetch content from volume with error {%s} for local file {%s}",
                 response.getStatusLine().getStatusCode(), localFilePath));
         status = VolumeOperationStatus.FAILED;
         errorMessage = "Failed to download file";
         return;
       }
-      HttpEntity entity = response.getEntity();
+      entity = response.getEntity();
       if (entity != null) {
         // Get the content of the HttpEntity
         InputStream inputStream = entity.getContent();
@@ -200,34 +248,31 @@ private void executePutOperation() {
     HttpPut httpPut = new HttpPut(operationUrl);
     headers.forEach(httpPut::addHeader);
 
-    // Set the FileEntity as the request body
-    File file = new File(localFilePath);
-    if (!file.exists() || file.isDirectory()) {
-      LoggingUtil.log(
-          LogLevel.ERROR,
-          String.format("Local file does not exist or is a directory {%s}", localFilePath));
-      status = VolumeOperationStatus.ABORTED;
-      errorMessage = "Local file does not exist or is a directory";
-      return;
-    }
-    if (file.length() == 0) {
-
-      LoggingUtil.log(LogLevel.ERROR, String.format("Local file is empty {%s}", localFilePath));
-      status = VolumeOperationStatus.ABORTED;
-      errorMessage = "Local file is empty";
-      return;
-    }
+    try {
+      if (statement.isAllowedInputStreamForVolumeOperation()) {
+        InputStreamEntity inputStream = statement.getInputStreamForUCVolume();
+        if (inputStream == null) {
+          status = VolumeOperationStatus.ABORTED;
+          errorMessage = "InputStream not set for PUT operation";
+          LoggingUtil.log(LogLevel.ERROR, errorMessage);
+          return;
+        }
+        httpPut.setEntity(inputStream);
+      } else {
+        // Set the FileEntity as the request body
+        File file = new File(localFilePath);
 
-    if (file.length() > PUT_SIZE_LIMITS) {
-      LoggingUtil.log(LogLevel.ERROR, String.format("Local file too large {%s}", localFilePath));
+        if (localFileHasErrorForPutOperation(file)) {
+          return;
+        }
+        httpPut.setEntity(new FileEntity(file, ContentType.DEFAULT_BINARY));
+      }
+    } catch (DatabricksSQLException e) {
       status = VolumeOperationStatus.ABORTED;
-      errorMessage = "Local file too large";
-      return;
+      errorMessage = "PUT operation called on closed statement";
+      LoggingUtil.log(LogLevel.ERROR, errorMessage);
     }
 
-    FileEntity fileEntity = new FileEntity(file, ContentType.DEFAULT_BINARY);
-    httpPut.setEntity(fileEntity);
-
     // Execute the request
     try (CloseableHttpResponse response = databricksHttpClient.execute(httpPut)) {
       // Process the response
@@ -254,6 +299,31 @@ private void executePutOperation() {
     }
   }
 
+  private boolean localFileHasErrorForPutOperation(File file) {
+    if (!file.exists() || file.isDirectory()) {
+      LoggingUtil.log(
+          LogLevel.ERROR,
+          String.format("Local file does not exist or is a directory {%s}", localFilePath));
+      status = VolumeOperationStatus.ABORTED;
+      errorMessage = "Local file does not exist or is a directory";
+      return true;
+    }
+    if (file.length() == 0) {
+      LoggingUtil.log(LogLevel.ERROR, String.format("Local file is empty {%s}", localFilePath));
+      status = VolumeOperationStatus.ABORTED;
+      errorMessage = "Local file is empty";
+      return true;
+    }
+
+    if (file.length() > PUT_SIZE_LIMITS) {
+      LoggingUtil.log(LogLevel.ERROR, String.format("Local file too large {%s}", localFilePath));
+      status = VolumeOperationStatus.ABORTED;
+      errorMessage = "Local file too large";
+      return true;
+    }
+    return false;
+  }
+
   private void executeDeleteOperation() {
     // TODO: Check for AWS specific handling
     HttpDelete httpDelete = new HttpDelete(operationUrl);

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

@@ -21,6 +21,8 @@ class VolumeOperationResult implements IExecutionResult {
   private final IDatabricksSession session;
   private final String statementId;
   private final IExecutionResult resultHandler;
+  private final IDatabricksResultSet resultSet;
+  private final IDatabricksStatement statement;
   private final IDatabricksHttpClient httpClient;
   private final long rowCount;
   private final long columnCount;
@@ -33,12 +35,16 @@ class VolumeOperationResult implements IExecutionResult {
       long totalRows,
       long totalColumns,
       IDatabricksSession session,
-      IExecutionResult resultHandler) {
+      IExecutionResult resultHandler,
+      IDatabricksStatement statement,
+      IDatabricksResultSet resultSet) {
     this.statementId = statementId;
     this.rowCount = totalRows;
     this.columnCount = totalColumns;
     this.session = session;
     this.resultHandler = resultHandler;
+    this.statement = statement;
+    this.resultSet = resultSet;
     this.httpClient = DatabricksHttpClient.getInstance(session.getConnectionContext());
     this.currentRowIndex = -1;
   }
@@ -49,12 +55,16 @@ class VolumeOperationResult implements IExecutionResult {
       ResultManifest manifest,
       IDatabricksSession session,
       IExecutionResult resultHandler,
-      IDatabricksHttpClient httpClient) {
+      IDatabricksHttpClient httpClient,
+      IDatabricksStatement statement,
+      IDatabricksResultSet resultSet) {
     this.statementId = statementId;
     this.rowCount = manifest.getTotalRowCount();
     this.columnCount = manifest.getSchema().getColumnCount();
     this.session = session;
     this.resultHandler = resultHandler;
+    this.statement = statement;
+    this.resultSet = resultSet;
     this.httpClient = httpClient;
     this.currentRowIndex = -1;
   }
@@ -73,7 +83,9 @@ private void initHandler(IExecutionResult resultHandler) throws DatabricksSQLExc
             session
                 .getClientInfoProperties()
                 .getOrDefault(ALLOWED_VOLUME_INGESTION_PATHS.toLowerCase(), ""),
-            httpClient);
+            httpClient,
+            statement,
+            resultSet);
     Thread thread = new Thread(volumeOperationExecutor);
     thread.setName("VolumeOperationExecutor " + statementId);
     thread.start();