(kyuubi) branch master updated: [KYUUBI #6639] Port HIVE-27815: Support update numModifiedRows

[chengpan]

This is an automated email from the ASF dual-hosted git repository.

chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 2d883e7cac [KYUUBI #6639] Port HIVE-27815: Support update 
numModifiedRows
2d883e7cac is described below

commit 2d883e7cac6e72f7792dabf69a2bb099df15b210
Author: Cheng Pan <cheng...@apache.org>
AuthorDate: Mon Aug 26 19:16:53 2024 +0800

    [KYUUBI #6639] Port HIVE-27815: Support update numModifiedRows
    
    # :mag: Description
    
    Backport https://github.com/apache/hive/pull/4819
    
    Note: it's only the JDBC driver side change, to make it work, we also need 
to modify the engines.
    
    ## Types of changes :bookmark:
    
    - [ ] Bugfix (non-breaking change which fixes an issue)
    - [x] New feature (non-breaking change which adds functionality)
    - [ ] Breaking change (fix or feature that would cause existing 
functionality to change)
    
    ## Test Plan 🧪
    
    We can not test this feature so far because we don't have engine support 
that. Pass GHA to ensure it breaks nothing.
    
    ---
    
    # Checklist 📝
    
    - [x] This patch was not authored or co-authored using [Generative 
Tooling](https://www.apache.org/legal/generative-tooling.html)
    
    **Be nice. Be informative.**
    
    Closes #6639 from pan3793/HIVE-27815.
    
    Closes #6639
    
    4b3bcd65f [Cheng Pan] fix
    c16dc28e8 [Cheng Pan] Port HIVE-27815: Support update numModifiedRows
    
    Authored-by: Cheng Pan <cheng...@apache.org>
    Signed-off-by: Cheng Pan <cheng...@apache.org>
---
 .../org/apache/kyuubi/jdbc/hive/KyuubiStatement.java   | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git 
a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java
 
b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java
index 23f0a1f43d..84965df203 100644
--- 
a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java
+++ 
b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java
@@ -380,7 +380,7 @@ public class KyuubiStatement implements SQLStatement, 
KyuubiLoggable {
     TGetOperationStatusResp statusResp = null;
 
     // Poll on the operation status, till the operation is complete
-    while (!isOperationComplete) {
+    do {
       try {
         /**
          * For an async SQLOperation, GetOperationStatus will use the long 
polling approach It will
@@ -427,7 +427,7 @@ public class KyuubiStatement implements SQLStatement, 
KyuubiLoggable {
         isLogBeingGenerated = false;
         throw new KyuubiSQLException(e.toString(), "08S01", e);
       }
-    }
+    } while (!isOperationComplete);
 
     /*
       we set progress bar to be completed when hive query execution has 
completed
@@ -513,7 +513,7 @@ public class KyuubiStatement implements SQLStatement, 
KyuubiLoggable {
   @Override
   public int executeUpdate(String sql) throws SQLException {
     execute(sql);
-    return 0;
+    return getUpdateCount();
   }
 
   @Override
@@ -589,8 +589,16 @@ public class KyuubiStatement implements SQLStatement, 
KyuubiLoggable {
      * client might end up using executeAsync and then call this to check if 
the query run is
      * finished.
      */
-    waitForOperationToComplete();
-    return -1;
+    long numModifiedRows = -1L;
+    TGetOperationStatusResp resp = waitForOperationToComplete();
+    if (resp != null && resp.isSetNumModifiedRows()) {
+      numModifiedRows = resp.getNumModifiedRows();
+    }
+    if (numModifiedRows == -1L || numModifiedRows > Integer.MAX_VALUE) {
+      LOG.warn("Invalid number of updated rows: {}", numModifiedRows);
+      return -1;
+    }
+    return (int) numModifiedRows;
   }
 
   @Override