Skip to content

Commit 36b95d3

Browse files
committed
[KYUUBI #1714] Add executeScala api for KyuubiStatement
<!-- Thanks for sending a pull request! Here are some tips for you: 1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html 2. If the PR is related to an issue in https://github.com/apache/incubator-kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'. 3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'. --> ### _Why are the changes needed?_ <!-- Please clarify why the changes are needed. For instance, 1. If you add a feature, you can talk about the use case of it. 2. If you fix a bug, you can clarify why it is a bug. --> Add executeScala api for KyuubiStatement, so that user can invoke this method to execute scala code directly. ### _How was this patch tested?_ - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request Closes #1714 from turboFei/execute_scala. Closes #1714 a62616d [Fei Wang] Add executeScala api for KyuubiStatement Authored-by: Fei Wang <fwang12@ebay.com> Signed-off-by: Fei Wang <fwang12@ebay.com>
1 parent 84839e0 commit 36b95d3

File tree

2 files changed

+42
-2
lines changed

2 files changed

+42
-2
lines changed

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/jdbc/KyuubiHiveDriverSuite.scala

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import java.util.Properties
2222
import org.apache.kyuubi.IcebergSuiteMixin
2323
import org.apache.kyuubi.engine.spark.WithSparkSQLEngine
2424
import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
25+
import org.apache.kyuubi.jdbc.hive.KyuubiStatement
2526
import org.apache.kyuubi.tags.IcebergTest
2627

2728
@IcebergTest
@@ -86,4 +87,19 @@ class KyuubiHiveDriverSuite extends WithSparkSQLEngine with IcebergSuiteMixin {
8687
connection.close()
8788
}
8889
}
90+
91+
test("add executeScala api for KyuubiStatement") {
92+
val driver = new KyuubiHiveDriver()
93+
val connection = driver.connect(getJdbcUrl, new Properties())
94+
val statement = connection.createStatement().asInstanceOf[KyuubiStatement]
95+
try {
96+
val code = """spark.sql("set kyuubi.operation.language").show(false)"""
97+
val resultSet = statement.executeScala(code)
98+
assert(resultSet.next())
99+
assert(resultSet.getString(1).contains("kyuubi.operation.language"))
100+
} finally {
101+
statement.close()
102+
connection.close()
103+
}
104+
}
89105
}

kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.sql.SQLTimeoutException;
2525
import java.sql.SQLWarning;
2626
import java.util.ArrayList;
27+
import java.util.Collections;
2728
import java.util.HashMap;
2829
import java.util.List;
2930
import java.util.Map;
@@ -245,7 +246,12 @@ public void closeOnCompletion() throws SQLException {
245246

246247
@Override
247248
public boolean execute(String sql) throws SQLException {
248-
runAsyncOnServer(sql);
249+
return executeWithConfOverlay(sql, null);
250+
}
251+
252+
private boolean executeWithConfOverlay(String sql, Map<String, String> confOverlay)
253+
throws SQLException {
254+
runAsyncOnServer(sql, confOverlay);
249255
TGetOperationStatusResp status = waitForOperationToComplete();
250256

251257
// The query should be completed by now
@@ -297,6 +303,10 @@ public boolean executeAsync(String sql) throws SQLException {
297303
}
298304

299305
private void runAsyncOnServer(String sql) throws SQLException {
306+
runAsyncOnServer(sql, null);
307+
}
308+
309+
private void runAsyncOnServer(String sql, Map<String, String> confOneTime) throws SQLException {
300310
checkConnection("execute");
301311

302312
closeClientOperation();
@@ -309,7 +319,13 @@ private void runAsyncOnServer(String sql) throws SQLException {
309319
* execution run asynchronously
310320
*/
311321
execReq.setRunAsync(true);
312-
execReq.setConfOverlay(sessConf);
322+
if (confOneTime != null) {
323+
Map<String, String> confOverlay = new HashMap<String, String>(sessConf);
324+
confOverlay.putAll(confOneTime);
325+
execReq.setConfOverlay(confOverlay);
326+
} else {
327+
execReq.setConfOverlay(sessConf);
328+
}
313329
execReq.setQueryTimeout(queryTimeout);
314330
try {
315331
TExecuteStatementResp execResp = client.ExecuteStatement(execReq);
@@ -483,6 +499,14 @@ public ResultSet executeQuery(String sql) throws SQLException {
483499
return resultSet;
484500
}
485501

502+
public ResultSet executeScala(String code) throws SQLException {
503+
if (!executeWithConfOverlay(
504+
code, Collections.singletonMap("kyuubi.operation.language", "SCALA"))) {
505+
throw new SQLException("The query did not generate a result set!");
506+
}
507+
return resultSet;
508+
}
509+
486510
/*
487511
* (non-Javadoc)
488512
*

0 commit comments

Comments
 (0)