[incubator-kyuubi] branch master updated: [KYUUBI #2451] Support isWrapperFor and unwrap

[ch...@apache.org]

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/incubator-kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 618732148 [KYUUBI #2451] Support isWrapperFor and unwrap
618732148 is described below

commit 618732148f546e3489a876d17e708e06a9f48c27
Author: gabrywu <ga...@apache.org>
AuthorDate: Tue Apr 26 09:56:55 2022 +0800

    [KYUUBI #2451] Support isWrapperFor and unwrap
    
    ### _Why are the changes needed?_
    support `isWrapperFor` and `unwrap` to make KyuubiConnection easy to use
    
    ### _How was this patch tested?_
    - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [x] Add screenshots for manual tests if appropriate
    
    - [x] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #2451 from gabrywu/firstPR.
    
    Closes #2451
    
    7eb63e3b [gabrywu] import org.apache.kyuubi.jdbc.hive.KyuubiConnection
    8078d4f0 [gabrywu] fix format violations
    3c236f12 [gabrywu] add unit tests
    a1b7d1a5 [gabrywu] add logic to isWrapperFor & unwrap
    
    Authored-by: gabrywu <ga...@apache.org>
    Signed-off-by: Cheng Pan <ch...@apache.org>
---
 .../scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala  | 11 ++++++++++-
 .../java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java    | 10 ++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala
index 26cd252b2..bf35358a5 100644
--- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala
@@ -22,7 +22,7 @@ import java.util.Properties
 import org.apache.kyuubi.IcebergSuiteMixin
 import org.apache.kyuubi.engine.spark.WithSparkSQLEngine
 import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
-import org.apache.kyuubi.jdbc.hive.KyuubiStatement
+import org.apache.kyuubi.jdbc.hive.{KyuubiConnection, KyuubiStatement}
 import org.apache.kyuubi.tags.IcebergTest
 
 @IcebergTest
@@ -102,4 +102,13 @@ class KyuubiHiveDriverSuite extends WithSparkSQLEngine with IcebergSuiteMixin {
       connection.close()
     }
   }
+
+  test("wrapable KyuubiConnection") {
+    val driver = new KyuubiHiveDriver()
+    val connection = driver.connect(getJdbcUrl, new Properties())
+    assert(connection.isWrapperFor(classOf[KyuubiConnection]))
+    val kyuubiConnection = connection.unwrap(classOf[KyuubiConnection])
+    assert(kyuubiConnection != null)
+    connection.close()
+  }
 }
diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java
index 1abc34168..c59e017cf 100644
--- a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java
+++ b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java
@@ -1585,8 +1585,7 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
 
   @Override
   public boolean isWrapperFor(Class<?> iface) throws SQLException {
-    // TODO Auto-generated method stub
-    throw new SQLFeatureNotSupportedException("Method not supported");
+    return iface.isInstance(this);
   }
 
   /*
@@ -1597,8 +1596,11 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
 
   @Override
   public <T> T unwrap(Class<T> iface) throws SQLException {
-    // TODO Auto-generated method stub
-    throw new SQLFeatureNotSupportedException("Method not supported");
+    if (!isWrapperFor(iface)) {
+      throw new SQLException(
+          this.getClass().getName() + " not unwrappable from " + iface.getName());
+    }
+    return iface.cast(this);
   }
 
   public TProtocolVersion getProtocol() {