[ES-1006739] Add content length parameter for input stream commands (#386)

* Adding GET extended command for UC volume operation

* Add changes for content length for input commands

Author: gopalldb

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

@@ -3,6 +3,7 @@
 import java.io.InputStream;
 import java.sql.SQLException;
 import java.util.List;
+import org.apache.http.entity.InputStreamEntity;
 
 public interface IDatabricksUCVolumeClient {
 
@@ -88,9 +89,9 @@ boolean getObject(
    * @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
+   * @return an instance of input stream entity
    */
-  InputStream getObject(String catalog, String schema, String volume, String objectPath)
+  InputStreamEntity getObject(String catalog, String schema, String volume, String objectPath)
       throws SQLException;
 
   /**
@@ -123,6 +124,7 @@ boolean putObject(
    * @param objectPath the destination path where the object (file) is to be uploaded from the
    *     volume as the root directory
    * @param inputStream the input stream from where the data is to be uploaded
+   * @param contentLength the length of the input stream
    * @param toOverwrite a boolean indicating whether to overwrite the object if it already exists
    * @return a boolean value indicating status of the PUT operation
    */
@@ -132,6 +134,7 @@ boolean putObject(
       String volume,
       String objectPath,
       InputStream inputStream,
+      long contentLength,
       boolean toOverwrite)
       throws SQLException;
 

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

@@ -12,6 +12,7 @@
 import java.sql.*;
 import java.util.ArrayList;
 import java.util.List;
+import org.apache.http.entity.InputStreamEntity;
 
 /** Implementation for DatabricksUCVolumeClient */
 public class DatabricksUCVolumeClient implements IDatabricksUCVolumeClient {
@@ -301,8 +302,8 @@ public boolean getObject(
   }
 
   @Override
-  public InputStream getObject(String catalog, String schema, String volume, String objectPath)
-      throws SQLException {
+  public InputStreamEntity getObject(
+      String catalog, String schema, String volume, String objectPath) throws SQLException {
 
     LoggingUtil.log(
         LogLevel.DEBUG,
@@ -375,6 +376,7 @@ public boolean putObject(
       String volume,
       String objectPath,
       InputStream inputStream,
+      long contentLength,
       boolean toOverwrite)
       throws SQLException {
 
@@ -392,7 +394,8 @@ public boolean putObject(
     try (Statement statement = connection.createStatement()) {
       IDatabricksStatement databricksStatement = (IDatabricksStatement) statement;
       databricksStatement.allowInputStreamForVolumeOperation(true);
-      databricksStatement.setInputStreamForUCVolume(inputStream);
+      databricksStatement.setInputStreamForUCVolume(
+          new InputStreamEntity(inputStream, contentLength));
 
       ResultSet resultSet = statement.executeQuery(putObjectQueryForInputStream);
       LoggingUtil.log(LogLevel.INFO, "PUT query executed successfully");

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

@@ -27,6 +27,7 @@
 import java.util.List;
 import java.util.Map;
 import org.apache.http.HttpEntity;
+import org.apache.http.entity.InputStreamEntity;
 
 public class DatabricksResultSet implements ResultSet, IDatabricksResultSet {
   protected static final String AFFECTED_ROWS_COUNT = "num_affected_rows";
@@ -41,6 +42,7 @@ public class DatabricksResultSet implements ResultSet, IDatabricksResultSet {
   private SQLWarning warnings = null;
   private boolean wasNull;
   private VolumeInputStream volumeInputStream = null;
+  private long volumeStreamContentLength = -1L;
 
   public DatabricksResultSet(
       StatementStatus statementStatus,
@@ -1623,12 +1625,13 @@ public void setVolumeOperationEntityStream(HttpEntity httpEntity)
     checkIfClosed();
     this.volumeInputStream =
         new VolumeInputStream(httpEntity, executionResult, this.parentStatement);
+    this.volumeStreamContentLength = httpEntity.getContentLength();
   }
 
   @Override
-  public InputStream getVolumeOperationInputStream() throws SQLException {
+  public InputStreamEntity getVolumeOperationInputStream() throws SQLException {
     checkIfClosed();
-    return this.volumeInputStream;
+    return new InputStreamEntity(this.volumeInputStream, this.volumeStreamContentLength);
   }
 
   private Object getObjectInternal(int columnIndex) throws SQLException {

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

@@ -10,11 +10,11 @@
 import com.databricks.jdbc.commons.LogLevel;
 import com.databricks.jdbc.commons.util.*;
 import com.google.common.annotations.VisibleForTesting;
-import java.io.InputStream;
 import java.sql.*;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.*;
+import org.apache.http.entity.InputStreamEntity;
 
 public class DatabricksStatement implements IDatabricksStatement, Statement {
 
@@ -27,7 +27,7 @@ public class DatabricksStatement implements IDatabricksStatement, Statement {
   private SQLWarning warnings = null;
   private int maxRows = DEFAULT_ROW_LIMIT;
   private boolean escapeProcessing = DEFAULT_ESCAPE_PROCESSING;
-  private InputStream inputStream = null;
+  private InputStreamEntity inputStream = null;
   private boolean allowInputStreamForUCVolume = false;
 
   public DatabricksStatement(DatabricksConnection connection) {
@@ -657,7 +657,8 @@ public boolean isAllowedInputStreamForVolumeOperation() throws DatabricksSQLExce
   }
 
   @Override
-  public void setInputStreamForUCVolume(InputStream inputStream) throws DatabricksSQLException {
+  public void setInputStreamForUCVolume(InputStreamEntity inputStream)
+      throws DatabricksSQLException {
     if (isAllowedInputStreamForVolumeOperation()) {
       this.inputStream = inputStream;
     } else {
@@ -666,7 +667,7 @@ public void setInputStreamForUCVolume(InputStream inputStream) throws Databricks
   }
 
   @Override
-  public InputStream getInputStreamForUCVolume() throws DatabricksSQLException {
+  public InputStreamEntity getInputStreamForUCVolume() throws DatabricksSQLException {
     if (isAllowedInputStreamForVolumeOperation()) {
       return inputStream;
     }

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

@@ -2,9 +2,9 @@
 
 import com.databricks.sdk.service.sql.StatementStatus;
 import java.io.IOException;
-import java.io.InputStream;
 import java.sql.SQLException;
 import org.apache.http.HttpEntity;
+import org.apache.http.entity.InputStreamEntity;
 
 public interface IDatabricksResultSet {
   String statementId();
@@ -17,5 +17,5 @@ public interface IDatabricksResultSet {
 
   void setVolumeOperationEntityStream(HttpEntity httpEntity) throws SQLException, IOException;
 
-  InputStream getVolumeOperationInputStream() throws SQLException;
+  InputStreamEntity getVolumeOperationInputStream() throws SQLException;
 }

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

@@ -1,8 +1,8 @@
 package com.databricks.jdbc.core;
 
-import java.io.InputStream;
 import java.sql.SQLException;
 import java.sql.Statement;
+import org.apache.http.entity.InputStreamEntity;
 
 /** Interface for Databricks specific statement. */
 public interface IDatabricksStatement {
@@ -26,7 +26,7 @@ public interface IDatabricksStatement {
 
   boolean isAllowedInputStreamForVolumeOperation() throws DatabricksSQLException;
 
-  void setInputStreamForUCVolume(InputStream inputStream) throws DatabricksSQLException;
+  void setInputStreamForUCVolume(InputStreamEntity inputStream) throws DatabricksSQLException;
 
-  InputStream getInputStreamForUCVolume() throws DatabricksSQLException;
+  InputStreamEntity getInputStreamForUCVolume() throws DatabricksSQLException;
 }

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

@@ -202,7 +202,6 @@ void testGetFloat() throws SQLException {
   void testGetUnicode() throws SQLException, UnsupportedEncodingException {
     DatabricksResultSet resultSet = getResultSet(StatementState.SUCCEEDED, null);
     String testString = "test";
-    InputStream expectedStream = new java.io.ByteArrayInputStream(testString.getBytes("UTF-16"));
     when(mockedExecutionResult.getObject(0)).thenReturn(testString);
     when(mockedResultSetMetadata.getColumnType(1)).thenReturn(Types.VARCHAR);
     assertNotNull(resultSet.getUnicodeStream(1));
@@ -716,8 +715,10 @@ void testVolumeOperationInputStream() throws Exception {
     DatabricksResultSet resultSet =
         getResultSet(StatementState.SUCCEEDED, mockedDatabricksStatement);
     HttpEntity mockEntity = mock(HttpEntity.class);
+    when(mockEntity.getContentLength()).thenReturn(10L);
     resultSet.setVolumeOperationEntityStream(mockEntity);
     assertNotNull(resultSet.getVolumeOperationInputStream());
+    assertEquals(10L, resultSet.getVolumeOperationInputStream().getContentLength());
   }
 
   @Test

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

@@ -20,6 +20,7 @@
 import java.util.Properties;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import org.apache.http.entity.InputStreamEntity;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -271,18 +272,20 @@ public void testInputStreamForVolumeOperation() throws Exception {
     assertFalse(statement.isAllowedInputStreamForVolumeOperation());
     assertNull(statement.getInputStreamForUCVolume());
     assertThrows(
-        DatabricksSQLException.class, () -> statement.setInputStreamForUCVolume(mockStream));
+        DatabricksSQLException.class,
+        () -> statement.setInputStreamForUCVolume(new InputStreamEntity(mockStream, -1L)));
 
     statement.allowInputStreamForVolumeOperation(true);
-    statement.setInputStreamForUCVolume(mockStream);
+    statement.setInputStreamForUCVolume(new InputStreamEntity(mockStream));
 
     assertTrue(statement.isAllowedInputStreamForVolumeOperation());
     assertNotNull(statement.getInputStreamForUCVolume());
 
     statement.close();
     assertThrows(DatabricksSQLException.class, statement::getInputStreamForUCVolume);
     assertThrows(
-        DatabricksSQLException.class, () -> statement.setInputStreamForUCVolume(mockStream));
+        DatabricksSQLException.class,
+        () -> statement.setInputStreamForUCVolume(new InputStreamEntity(mockStream, -1L)));
     assertThrows(DatabricksSQLException.class, statement::isAllowedInputStreamForVolumeOperation);
     assertThrows(
         DatabricksSQLException.class, () -> statement.allowInputStreamForVolumeOperation(false));

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

@@ -15,6 +15,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Stream;
+import org.apache.http.entity.InputStreamEntity;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -527,7 +528,7 @@ private static Stream<Arguments> provideParametersForDeleteObject_InvalidObjectP
   @ParameterizedTest
   @MethodSource("provideParametersForGetObjectWithInputStream")
   public void testGetObjectWithInputStream(
-      String catalog, String schema, String volume, String objectPath, InputStream expected)
+      String catalog, String schema, String volume, String objectPath, InputStreamEntity expected)
       throws SQLException {
     DatabricksUCVolumeClient client = new DatabricksUCVolumeClient(connection);
 
@@ -541,16 +542,17 @@ public void testGetObjectWithInputStream(
     when(databricksResultSet.next()).thenReturn(true);
     when(databricksResultSet.getVolumeOperationInputStream()).thenReturn(expected);
 
-    InputStream result = client.getObject(catalog, schema, volume, objectPath);
+    InputStreamEntity 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));
+    InputStreamEntity inputStream =
+        new InputStreamEntity(
+            new ByteArrayInputStream("test data".getBytes(StandardCharsets.UTF_8)), 10L);
     return Stream.of(
         Arguments.of("test_catalog", "test_schema", "test_volume", "test_objectpath", inputStream));
   }
@@ -563,6 +565,7 @@ public void testPutObjectWithInputStream(
       String volume,
       String objectPath,
       InputStream inputStream,
+      long length,
       boolean toOverwrite,
       boolean expected)
       throws SQLException {
@@ -579,7 +582,7 @@ public void testPutObjectWithInputStream(
         .thenReturn(VOLUME_OPERATION_STATUS_SUCCEEDED);
 
     boolean result =
-        client.putObject(catalog, schema, volume, objectPath, inputStream, toOverwrite);
+        client.putObject(catalog, schema, volume, objectPath, inputStream, length, toOverwrite);
 
     assertEquals(expected, result);
     verify(databricksStatement).executeQuery(putObjectQuery);
@@ -596,6 +599,7 @@ private static Stream<Arguments> provideParametersForPutObjectWithInputStream()
             "test_volume",
             "test_objectpath",
             inputStream,
+            10L,
             false,
             true));
   }

src/test/java/com/databricks/jdbc/integration/e2e/UCVolumeTests.java

@@ -527,7 +527,7 @@ public void testPutObjectWithInputStream() throws Exception {
 
     boolean result =
         client.putObject(
-            UC_VOLUME_CATALOG, UC_VOLUME_SCHEMA, volume, objectPath, inputStream, toOverwrite);
+            UC_VOLUME_CATALOG, UC_VOLUME_SCHEMA, volume, objectPath, inputStream, 10L, toOverwrite);
     assertEquals(expected, result);
   }
 }