Merge branch 'main' into PECOBLR-849

Author: msrathore-db

NEXT_CHANGELOG.md

@@ -18,5 +18,6 @@
 ### Fixed
 - Integrated Azure U2M flow into driver for improved stability.
 - Fixed `ResultSet.getString` for Boolean columns in Metadata result set.
+- Fixed volume operations not completing unless the ResultSet is fully iterated.
 ---
 *Note: When making changes, please add your change under the appropriate section with a brief description.* 

src/main/java/com/databricks/jdbc/api/impl/volume/VolumeOperationResult.java

@@ -47,7 +47,8 @@ public VolumeOperationResult(
       long totalColumns,
       IDatabricksSession session,
       IExecutionResult resultHandler,
-      IDatabricksStatementInternal statement) {
+      IDatabricksStatementInternal statement)
+      throws DatabricksSQLException {
     this.rowCount = totalRows;
     this.columnCount = totalColumns;
     this.session = session;
@@ -57,6 +58,7 @@ public VolumeOperationResult(
         DatabricksHttpClientFactory.getInstance()
             .getClient(session.getConnectionContext(), HttpClientType.VOLUME);
     this.currentRowIndex = -1;
+    completeVolumeOperation();
   }
 
   @VisibleForTesting
@@ -65,14 +67,16 @@ public VolumeOperationResult(
       IDatabricksSession session,
       IExecutionResult resultHandler,
       IDatabricksHttpClient httpClient,
-      IDatabricksStatementInternal statement) {
+      IDatabricksStatementInternal statement)
+      throws DatabricksSQLException {
     this.rowCount = manifest.getTotalRowCount();
     this.columnCount = manifest.getSchema().getColumnCount();
     this.session = session;
     this.resultHandler = resultHandler;
     this.statement = statement;
     this.httpClient = httpClient;
     this.currentRowIndex = -1;
+    completeVolumeOperation();
   }
 
   private void initHandler(IExecutionResult resultHandler) throws DatabricksSQLException {
@@ -181,7 +185,12 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
     if (columnIndex == 0) {
       return volumeOperationProcessor.getStatus().name();
     }
-    String errorMessage = (currentRowIndex < 0) ? "Invalid row access" : "Invalid column access";
+    String userMessage = (currentRowIndex < 0) ? "Invalid row access" : "Invalid column access";
+    String errorMessage =
+        String.format(
+            "%s, Row: %s, Column: %s, statementID %s",
+            userMessage, currentRowIndex, columnIndex, statement.getStatementId());
+    LOGGER.error(errorMessage);
     throw new DatabricksVolumeOperationException(
         errorMessage, DatabricksDriverErrorCode.VOLUME_OPERATION_INVALID_STATE);
   }
@@ -193,16 +202,21 @@ public long getCurrentRow() {
 
   @Override
   public boolean next() throws DatabricksSQLException {
-    if (hasNext()) {
+    if (currentRowIndex == -1) {
+      currentRowIndex++;
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  private void completeVolumeOperation() throws DatabricksSQLException {
+    while (resultHandler.hasNext()) {
       validateMetadata();
       resultHandler.next();
       initHandler(resultHandler);
       volumeOperationProcessor.process();
       ensureSuccessVolumeProcessorStatus();
-      currentRowIndex++;
-      return true;
-    } else {
-      return false;
     }
   }
 
@@ -217,7 +231,7 @@ public InputStreamEntity getVolumeOperationInputStream() {
 
   @Override
   public boolean hasNext() {
-    return resultHandler.hasNext();
+    return currentRowIndex == -1;
   }
 
   @Override