From 2d883e7cac6e72f7792dabf69a2bb099df15b210 Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@apache.org>
Date: Mon, 26 Aug 2024 19:16:53 +0800
Subject: [PATCH] [KYUUBI #6639] Port HIVE-27815: Support update
 numModifiedRows
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# :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 <chengpan@apache.org>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../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 23f0a1f43..84965df20 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