You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2022/05/26 10:06:40 UTC
[incubator-kyuubi] branch master updated: [KYUUBI #2721] Implement dedicated set/get catalog/database operators
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 9b5023070 [KYUUBI #2721] Implement dedicated set/get catalog/database operators
9b5023070 is described below
commit 9b5023070aedba1065ad6e0feac62ab0739e6611
Author: sychen <sy...@ctrip.com>
AuthorDate: Thu May 26 18:06:30 2022 +0800
[KYUUBI #2721] Implement dedicated set/get catalog/database operators
### _Why are the changes needed?_
close #2721
### _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 #2728 from cxzl25/KYUUBI-2721.
Closes #2721
67051644 [sychen] use placeholders and confOverlay
2acf0a45 [sychen] unused import
8bc40b9f [sychen] fix UT
3495b280 [sychen] fix jdbc
e55bb2e2 [sychen] style
ff9078ab [sychen] fix UT
a0247c7e [sychen] Merge branch 'master' into KYUUBI-2721
f62461ce [sychen] use placeholders
698fbbdd [sychen] set catalog/schema empty resultset
45caf552 [sychen] fix import
7f13ad8d [sychen] fix server spark catalog UT
4982ec71 [sychen] hive engine 1.8
9877038f [sychen] add UT
21e57870 [sychen] set/get catalog/database operators
Authored-by: sychen <sy...@ctrip.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
docs/deployment/settings.md | 1 +
.../flink/operation/FlinkSQLOperationManager.scala | 31 ++++++++++
.../engine/flink/operation/GetCurrentCatalog.scala | 25 +++++---
.../flink/operation/GetCurrentDatabase.scala | 25 +++++---
.../engine/flink/operation/SetCurrentCatalog.scala | 23 +++++---
.../flink/operation/SetCurrentDatabase.scala | 23 +++++---
.../flink/operation/FlinkOperationSuite.scala | 38 +++++++++++-
.../engine/hive/operation/GetCurrentCatalog.scala | 27 ++++++---
.../engine/hive/operation/GetCurrentDatabase.scala | 27 ++++++---
.../hive/operation/HiveOperationManager.scala | 27 +++++++++
.../engine/hive/operation/SetCurrentCatalog.scala | 32 +++++++---
.../engine/hive/operation/SetCurrentDatabase.scala | 27 ++++++---
.../HiveCatalogDatabaseOperationSuite.scala | 68 +++++++++++++++++++++
.../jdbc/operation/JdbcOperationManager.scala | 16 +++++
.../engine/spark/operation/GetCurrentCatalog.scala | 30 +++++++---
.../spark/operation/GetCurrentDatabase.scala | 30 +++++++---
.../engine/spark/operation/SetCurrentCatalog.scala | 29 ++++++---
.../spark/operation/SetCurrentDatabase.scala | 29 ++++++---
.../spark/operation/SparkSQLOperationManager.scala | 29 +++++++++
.../engine/spark/shim/CatalogShim_v2_4.scala | 14 +++++
.../engine/spark/shim/CatalogShim_v3_0.scala | 21 +++++++
.../engine/spark/shim/SparkCatalogShim.scala | 8 +++
.../SparkCatalogDatabaseOperationSuite.scala | 69 ++++++++++++++++++++++
.../kyuubi/engine/trino/TrinoStatement.scala | 2 +
.../engine/trino/operation/GetCurrentCatalog.scala | 44 ++++++++++++++
.../trino/operation/GetCurrentDatabase.scala | 44 ++++++++++++++
.../engine/trino/operation/SetCurrentCatalog.scala | 27 ++++++---
.../trino/operation/SetCurrentDatabase.scala | 27 ++++++---
.../trino/operation/TrinoOperationManager.scala | 28 ++++++++-
.../TrinoCatalogDatabaseOperationSuite.scala | 55 +++++++++++++++++
.../org/apache/kyuubi/config/KyuubiConf.scala | 8 +++
.../apache/kyuubi/config/KyuubiReservedKeys.scala | 4 ++
.../apache/kyuubi/operation/OperationManager.scala | 31 ++++++++++
.../kyuubi/operation/NoopOperationManager.scala | 20 +++++++
.../apache/kyuubi/jdbc/hive/KyuubiConnection.java | 33 ++++++++---
.../apache/kyuubi/jdbc/hive/KyuubiStatement.java | 30 ++++++++++
.../kyuubi/operation/KyuubiOperationManager.scala | 17 ++++++
.../operation/KyuubiOperationManagerSuite.scala | 32 +++++++++-
38 files changed, 912 insertions(+), 139 deletions(-)
diff --git a/docs/deployment/settings.md b/docs/deployment/settings.md
index fee96f956..d1d88d972 100644
--- a/docs/deployment/settings.md
+++ b/docs/deployment/settings.md
@@ -223,6 +223,7 @@ Key | Default | Meaning | Type | Since
<code>kyuubi.engine.jdbc.connection.user</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'><undefined></div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>The user is used for connecting to server</div>|<div style='width: 30pt'>string</div>|<div style='width: 20pt'>1.6.0</div>
<code>kyuubi.engine.jdbc.driver.class</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'><undefined></div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>The driver class for jdbc engine connection</div>|<div style='width: 30pt'>string</div>|<div style='width: 20pt'>1.6.0</div>
<code>kyuubi.engine.jdbc.type</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'><undefined></div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>The short name of jdbc type</div>|<div style='width: 30pt'>string</div>|<div style='width: 20pt'>1.6.0</div>
+<code>kyuubi.engine.operation.convert.catalog.database.enabled</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'>true</div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines</div>|<div style='width: 30pt'>boolean</div>|<div style='width: 20pt'>1.6.0</div>
<code>kyuubi.engine.operation.log.dir.root</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'>engine_operation_logs</div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>Root directory for query operation log at engine-side.</div>|<div style='width: 30pt'>string</div>|<div style='width: 20pt'>1.4.0</div>
<code>kyuubi.engine.pool.name</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'>engine-pool</div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>The name of engine pool.</div>|<div style='width: 30pt'>string</div>|<div style='width: 20pt'>1.5.0</div>
<code>kyuubi.engine.pool.size</code>|<div style='width: 65pt;word-wrap: break-word;white-space: normal'>-1</div>|<div style='width: 170pt;word-wrap: break-word;white-space: normal'>The size of engine pool. Note that, if the size is less than 1, the engine pool will not be enabled; otherwise, the size of the engine pool will be min(this, kyuubi.engine.pool.size.threshold).</div>|<div style='width: 30pt'>int</div>|<div style='width: 20pt'>1.4.0</div>
diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala
index 9a98fe1b7..056e7844f 100644
--- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala
@@ -36,6 +36,9 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage
private lazy val resultMaxRowsDefault = getConf.get(ENGINE_FLINK_MAX_ROWS)
+ private lazy val operationConvertCatalogDatabaseDefault =
+ getConf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)
+
override def newExecuteStatementOperation(
session: Session,
statement: String,
@@ -43,6 +46,14 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage
runAsync: Boolean,
queryTimeout: Long): Operation = {
val flinkSession = session.asInstanceOf[FlinkSessionImpl]
+ if (flinkSession.sessionContext.getConfigMap.getOrDefault(
+ ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key,
+ operationConvertCatalogDatabaseDefault.toString).toBoolean) {
+ val catalogDatabaseOperation = processCatalogDatabase(session, statement, confOverlay)
+ if (catalogDatabaseOperation != null) {
+ return catalogDatabaseOperation
+ }
+ }
val mode = flinkSession.sessionContext.getConfigMap.getOrDefault(
OPERATION_PLAN_ONLY_MODE.key,
operationModeDefault)
@@ -59,6 +70,26 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage
addOperation(op)
}
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ val op = new SetCurrentCatalog(session, catalog)
+ addOperation(op)
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ val op = new GetCurrentCatalog(session)
+ addOperation(op)
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ val op = new SetCurrentDatabase(session, database)
+ addOperation(op)
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ val op = new GetCurrentDatabase(session)
+ addOperation(op)
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val op = new GetTypeInfo(session)
addOperation(op)
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala
similarity index 54%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala
index 39bef393d..2ddb06361 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala
@@ -15,14 +15,21 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.flink.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.kyuubi.engine.flink.result.ResultSetUtil
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant.TABLE_CAT
+import org.apache.kyuubi.session.Session
+
+class GetCurrentCatalog(session: Session)
+ extends FlinkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+ val catalog = tableEnv.getCurrentCatalog
+ resultSet = ResultSetUtil.stringListToResultSet(List(catalog), TABLE_CAT)
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala
similarity index 54%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala
index 39bef393d..1e392511e 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala
@@ -15,14 +15,21 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.flink.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.kyuubi.engine.flink.result.ResultSetUtil
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant.TABLE_SCHEM
+import org.apache.kyuubi.session.Session
+
+class GetCurrentDatabase(session: Session)
+ extends FlinkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+ val database = tableEnv.getCurrentDatabase
+ resultSet = ResultSetUtil.stringListToResultSet(List(database), TABLE_SCHEM)
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala
similarity index 61%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala
index 39bef393d..6090d5040 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala
@@ -15,14 +15,19 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.flink.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentCatalog(session: Session, catalog: String)
+ extends FlinkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+ tableEnv.useCatalog(catalog)
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala
similarity index 61%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala
index 39bef393d..633a11974 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala
@@ -15,14 +15,19 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.flink.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentDatabase(session: Session, database: String)
+ extends FlinkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+ tableEnv.useDatabase(database)
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperationSuite.scala b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperationSuite.scala
index f96cc61bc..90dd00257 100644
--- a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperationSuite.scala
+++ b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperationSuite.scala
@@ -26,7 +26,7 @@ import org.apache.hive.service.rpc.thrift.{TExecuteStatementReq, TFetchResultsRe
import org.scalatest.concurrent.PatienceConfiguration.Timeout
import org.scalatest.time.SpanSugar._
-import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf._
import org.apache.kyuubi.config.KyuubiConf.OperationModes.NONE
import org.apache.kyuubi.engine.flink.WithFlinkSQLEngine
import org.apache.kyuubi.engine.flink.result.Constants
@@ -37,7 +37,7 @@ import org.apache.kyuubi.service.ServiceState._
class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper {
override def withKyuubiConf: Map[String, String] =
- Map(KyuubiConf.OPERATION_PLAN_ONLY_MODE.key -> NONE.toString)
+ Map(OPERATION_PLAN_ONLY_MODE.key -> NONE.toString)
override protected def jdbcUrl: String =
s"jdbc:hive2://${engine.frontendServices.head.connectionUrl}/;"
@@ -665,6 +665,22 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper {
})
}
+ test("execute statement - set/get catalog") {
+ withSessionConf()(
+ Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true"))(
+ Map.empty) {
+ withJdbcStatement() { statement =>
+ statement.executeQuery("create catalog cat_a with ('type'='generic_in_memory')")
+ val catalog = statement.getConnection.getCatalog
+ assert(catalog == "default_catalog")
+ statement.getConnection.setCatalog("cat_a")
+ val changedCatalog = statement.getConnection.getCatalog
+ assert(changedCatalog == "cat_a")
+ assert(statement.execute("drop catalog cat_a"))
+ }
+ }
+ }
+
test("execute statement - create/alter/drop database") {
// TODO: validate table results after FLINK-25558 is resolved
withJdbcStatement()({ statement =>
@@ -674,6 +690,22 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper {
})
}
+ test("execute statement - set/get database") {
+ withSessionConf()(
+ Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true"))(
+ Map.empty) {
+ withJdbcStatement()({ statement =>
+ statement.executeQuery("create database db_a")
+ val schema = statement.getConnection.getSchema
+ assert(schema == "default_database")
+ statement.getConnection.setSchema("db_a")
+ val changedSchema = statement.getConnection.getSchema
+ assert(changedSchema == "db_a")
+ assert(statement.execute("drop database db_a"))
+ })
+ }
+ }
+
test("execute statement - create/alter/drop table") {
// TODO: validate table results after FLINK-25558 is resolved
withJdbcStatement()({ statement =>
@@ -777,7 +809,7 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper {
}
test("ensure result max rows") {
- withSessionConf()(Map(KyuubiConf.ENGINE_FLINK_MAX_ROWS.key -> "200"))(Map.empty) {
+ withSessionConf()(Map(ENGINE_FLINK_MAX_ROWS.key -> "200"))(Map.empty) {
withJdbcStatement() { statement =>
statement.execute("create table tbl_src (a bigint) with ('connector' = 'datagen')")
val resultSet = statement.executeQuery(s"select a from tbl_src")
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala
similarity index 57%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala
index 39bef393d..ee6d44df0 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala
@@ -15,14 +15,23 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.hive.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import scala.collection.JavaConverters._
+
+import org.apache.hive.service.cli.operation.Operation
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class GetCurrentCatalog(session: Session)
+ extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) {
+ // Hive does not support catalog
+ override val internalHiveOperation: Operation =
+ delegatedOperationManager.newExecuteStatementOperation(
+ hive,
+ "SELECT '' AS TABLE_CAT",
+ Map.empty[String, String].asJava,
+ false,
+ 0)
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala
similarity index 58%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala
index 39bef393d..3bdbffa15 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala
@@ -15,14 +15,23 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.hive.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import scala.collection.JavaConverters._
+
+import org.apache.hive.service.cli.operation.Operation
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class GetCurrentDatabase(session: Session)
+ extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override val internalHiveOperation: Operation =
+ delegatedOperationManager.newExecuteStatementOperation(
+ hive,
+ "SELECT current_database()",
+ Map.empty[String, String].asJava,
+ false,
+ 0)
}
diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
index cdfea54fa..8cd7a3112 100644
--- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
@@ -22,6 +22,7 @@ import java.util.List
import org.apache.hadoop.hive.conf.HiveConf.ConfVars
import org.apache.hive.service.rpc.thrift.TRowSet
+import org.apache.kyuubi.config.KyuubiConf._
import org.apache.kyuubi.operation.{Operation, OperationHandle, OperationManager}
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.session.Session
@@ -36,10 +37,36 @@ class HiveOperationManager() extends OperationManager("HiveOperationManager") {
confOverlay: Map[String, String],
runAsync: Boolean,
queryTimeout: Long): Operation = {
+ if (session.sessionManager.getConf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)) {
+ val catalogDatabaseOperation = processCatalogDatabase(session, statement, confOverlay)
+ if (catalogDatabaseOperation != null) {
+ return catalogDatabaseOperation
+ }
+ }
val operation = new ExecuteStatement(session, statement, confOverlay, runAsync, queryTimeout)
addOperation(operation)
}
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ val op = new SetCurrentCatalog(session, catalog)
+ addOperation(op)
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ val op = new GetCurrentCatalog(session)
+ addOperation(op)
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ val op = new SetCurrentDatabase(session, database)
+ addOperation(op)
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ val op = new GetCurrentDatabase(session)
+ addOperation(op)
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val operation = new GetTypeInfo(session)
addOperation(operation)
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala
similarity index 53%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala
index 39bef393d..24340fb2c 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala
@@ -15,14 +15,28 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.hive.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import scala.collection.JavaConverters._
+
+import org.apache.hive.service.cli.operation.Operation
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentCatalog(session: Session, catalog: String)
+ extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) {
+ // Hive does not support catalog
+ override val internalHiveOperation: Operation =
+ delegatedOperationManager.newExecuteStatementOperation(
+ hive,
+ "SELECT ''",
+ Map.empty[String, String].asJava,
+ false,
+ 0)
+
+ override def runInternal(): Unit = {
+ super.runInternal()
+ setHasResultSet(false)
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala
similarity index 58%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala
index 39bef393d..c23127d38 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala
@@ -15,14 +15,23 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.hive.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import scala.collection.JavaConverters._
+
+import org.apache.hive.service.cli.operation.Operation
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentDatabase(session: Session, database: String)
+ extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override val internalHiveOperation: Operation =
+ delegatedOperationManager.newExecuteStatementOperation(
+ hive,
+ "USE " + database,
+ Map.empty[String, String].asJava,
+ false,
+ 0)
}
diff --git a/externals/kyuubi-hive-sql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala b/externals/kyuubi-hive-sql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala
new file mode 100644
index 000000000..aa95c4df6
--- /dev/null
+++ b/externals/kyuubi-hive-sql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.hive.operation
+
+import org.apache.commons.lang3.{JavaVersion, SystemUtils}
+
+import org.apache.kyuubi.Utils
+import org.apache.kyuubi.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED
+import org.apache.kyuubi.engine.hive.HiveSQLEngine
+import org.apache.kyuubi.operation.HiveJDBCTestHelper
+
+class HiveCatalogDatabaseOperationSuite extends HiveJDBCTestHelper {
+
+ override def beforeAll(): Unit = {
+ val metastore = Utils.createTempDir(namePrefix = getClass.getSimpleName)
+ metastore.toFile.delete()
+ val args = Array(
+ "--conf",
+ s"javax.jdo.option.ConnectionURL=jdbc:derby:;databaseName=$metastore;create=true",
+ "--conf",
+ s"${ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key}=true")
+ HiveSQLEngine.main(args)
+ super.beforeAll()
+ }
+
+ override protected def jdbcUrl: String = {
+ "jdbc:hive2://" + HiveSQLEngine.currentEngine.get.frontendServices.head.connectionUrl + "/;"
+ }
+
+ test("test set/get catalog") {
+ assume(SystemUtils.isJavaVersionAtMost(JavaVersion.JAVA_1_8))
+ withJdbcStatement()({ statement =>
+ val catalog = statement.getConnection.getCatalog
+ assert(catalog == "")
+ statement.getConnection.setCatalog("hive_not_support_catalog")
+ val changedCatalog = statement.getConnection.getCatalog
+ assert(changedCatalog == "")
+ })
+ }
+
+ test("test set/get database") {
+ assume(SystemUtils.isJavaVersionAtMost(JavaVersion.JAVA_1_8))
+ withJdbcStatement()({ statement =>
+ statement.execute("create database test_hive_db")
+ val schema = statement.getConnection.getSchema
+ assert(schema == "default")
+ statement.getConnection.setSchema("test_hive_db")
+ val changedSchema = statement.getConnection.getSchema
+ assert(changedSchema == "test_hive_db")
+ statement.execute("drop database test_hive_db")
+ })
+ }
+}
diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperationManager.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperationManager.scala
index d4455e93c..c9700870b 100644
--- a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperationManager.scala
+++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperationManager.scala
@@ -122,4 +122,20 @@ class JdbcOperationManager(conf: KyuubiConf) extends OperationManager("JdbcOpera
override def getQueryId(operation: Operation): String = {
throw KyuubiSQLException.featureNotSupported()
}
+
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ throw KyuubiSQLException.featureNotSupported()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala
similarity index 50%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala
index 39bef393d..47b386016 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala
@@ -15,14 +15,26 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.spark.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.spark.sql.types.StructType
+
+import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
+import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType}
+import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant.TABLE_CAT
+import org.apache.kyuubi.session.Session
+
+class GetCurrentCatalog(session: Session)
+ extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def resultSchema: StructType = {
+ new StructType()
+ .add(TABLE_CAT, "string", nullable = true, "Catalog name.")
+ }
+
+ override protected def runInternal(): Unit = {
+ try {
+ iter = new IterableFetchIterator(Seq(SparkCatalogShim().getCurrentCatalog(spark)))
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala
similarity index 50%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala
index 39bef393d..954e23898 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala
@@ -15,14 +15,26 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.spark.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.spark.sql.types.StructType
+
+import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
+import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType}
+import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant.TABLE_SCHEM
+import org.apache.kyuubi.session.Session
+
+class GetCurrentDatabase(session: Session)
+ extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def resultSchema: StructType = {
+ new StructType()
+ .add(TABLE_SCHEM, "string", nullable = true, "Schema name.")
+ }
+
+ override protected def runInternal(): Unit = {
+ try {
+ iter = new IterableFetchIterator(Seq(SparkCatalogShim().getCurrentDatabase(spark)))
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala
similarity index 56%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala
index 39bef393d..dd0bf987c 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala
@@ -15,14 +15,25 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.spark.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.spark.sql.types.StructType
+
+import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentCatalog(session: Session, catalog: String)
+ extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def resultSchema: StructType = {
+ new StructType()
+ }
+
+ override protected def runInternal(): Unit = {
+ try {
+ SparkCatalogShim().setCurrentCatalog(spark, catalog)
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala
similarity index 56%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala
index 39bef393d..b4efc15b0 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala
@@ -15,14 +15,25 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.spark.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import org.apache.spark.sql.types.StructType
+
+import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentDatabase(session: Session, database: String)
+ extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def resultSchema: StructType = {
+ new StructType()
+ }
+
+ override protected def runInternal(): Unit = {
+ try {
+ SparkCatalogShim().setCurrentDatabase(spark, database)
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala
index bc7b84390..30ad1c203 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala
@@ -38,6 +38,8 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n
private lazy val operationModeDefault = getConf.get(OPERATION_PLAN_ONLY_MODE)
private lazy val operationIncrementalCollectDefault = getConf.get(OPERATION_INCREMENTAL_COLLECT)
private lazy val operationLanguageDefault = getConf.get(OPERATION_LANGUAGE)
+ private lazy val operationConvertCatalogDatabaseDefault =
+ getConf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)
private val sessionToRepl = new ConcurrentHashMap[SessionHandle, KyuubiSparkILoop]().asScala
@@ -53,6 +55,13 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n
runAsync: Boolean,
queryTimeout: Long): Operation = {
val spark = session.asInstanceOf[SparkSessionImpl].spark
+ if (spark.conf.getOption(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key)
+ .map(_.toBoolean).getOrElse(operationConvertCatalogDatabaseDefault)) {
+ val catalogDatabaseOperation = processCatalogDatabase(session, statement, confOverlay)
+ if (catalogDatabaseOperation != null) {
+ return catalogDatabaseOperation
+ }
+ }
val lang = confOverlay.getOrElse(
OPERATION_LANGUAGE.key,
spark.conf.get(OPERATION_LANGUAGE.key, operationLanguageDefault))
@@ -75,6 +84,26 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n
addOperation(operation)
}
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ val op = new SetCurrentCatalog(session, catalog)
+ addOperation(op)
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ val op = new GetCurrentCatalog(session)
+ addOperation(op)
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ val op = new SetCurrentDatabase(session, database)
+ addOperation(op)
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ val op = new GetCurrentDatabase(session)
+ addOperation(op)
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val op = new GetTypeInfo(session)
addOperation(op)
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v2_4.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v2_4.scala
index cc5a25b0c..5977cd415 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v2_4.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v2_4.scala
@@ -30,6 +30,12 @@ class CatalogShim_v2_4 extends SparkCatalogShim {
override protected def catalogExists(spark: SparkSession, catalog: String): Boolean = false
+ override def setCurrentCatalog(spark: SparkSession, catalog: String): Unit = {}
+
+ override def getCurrentCatalog(spark: SparkSession): Row = {
+ Row(SparkCatalogShim.SESSION_CATALOG)
+ }
+
override def getSchemas(
spark: SparkSession,
catalogName: String,
@@ -38,6 +44,14 @@ class CatalogShim_v2_4 extends SparkCatalogShim {
getGlobalTempViewManager(spark, schemaPattern)).map(Row(_, ""))
}
+ def setCurrentDatabase(spark: SparkSession, databaseName: String): Unit = {
+ spark.sessionState.catalog.setCurrentDatabase(databaseName)
+ }
+
+ def getCurrentDatabase(spark: SparkSession): Row = {
+ Row(spark.sessionState.catalog.getCurrentDatabase)
+ }
+
override protected def getGlobalTempViewManager(
spark: SparkSession,
schemaPattern: String): Seq[String] = {
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v3_0.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v3_0.scala
index db0335849..e42d2cf33 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v3_0.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/CatalogShim_v3_0.scala
@@ -54,6 +54,19 @@ class CatalogShim_v3_0 extends CatalogShim_v2_4 {
spark.sessionState.catalogManager.isCatalogRegistered(catalog)
}
+ override def setCurrentCatalog(spark: SparkSession, catalog: String): Unit = {
+ // SPARK-36841(3.3.0) Ensure setCurrentCatalog method catalog must exist
+ if (spark.sessionState.catalogManager.isCatalogRegistered(catalog)) {
+ spark.sessionState.catalogManager.setCurrentCatalog(catalog)
+ } else {
+ throw new IllegalArgumentException(s"Cannot find catalog plugin class for catalog '$catalog'")
+ }
+ }
+
+ override def getCurrentCatalog(spark: SparkSession): Row = {
+ Row(spark.sessionState.catalogManager.currentCatalog.name())
+ }
+
private def listAllNamespaces(
catalog: SupportsNamespaces,
namespaces: Array[Array[String]]): Array[Array[String]] = {
@@ -122,6 +135,14 @@ class CatalogShim_v3_0 extends CatalogShim_v2_4 {
(schemas ++ viewMgr).map(Row(_, catalog.name()))
}
+ override def setCurrentDatabase(spark: SparkSession, databaseName: String): Unit = {
+ spark.sessionState.catalogManager.setCurrentNamespace(Array(databaseName))
+ }
+
+ override def getCurrentDatabase(spark: SparkSession): Row = {
+ Row(spark.sessionState.catalogManager.currentNamespace.map(quoteIfNeeded).mkString("."))
+ }
+
override def getCatalogTablesOrViews(
spark: SparkSession,
catalogName: String,
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/SparkCatalogShim.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/SparkCatalogShim.scala
index 831b39c74..bc5792823 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/SparkCatalogShim.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/shim/SparkCatalogShim.scala
@@ -41,6 +41,10 @@ trait SparkCatalogShim extends Logging {
protected def catalogExists(spark: SparkSession, catalog: String): Boolean
+ def setCurrentCatalog(spark: SparkSession, catalog: String): Unit
+
+ def getCurrentCatalog(spark: SparkSession): Row
+
// ///////////////////////////////////////////////////////////////////////////////////////////////
// Schema //
// ///////////////////////////////////////////////////////////////////////////////////////////////
@@ -50,6 +54,10 @@ trait SparkCatalogShim extends Logging {
*/
def getSchemas(spark: SparkSession, catalogName: String, schemaPattern: String): Seq[Row]
+ def setCurrentDatabase(spark: SparkSession, databaseName: String): Unit
+
+ def getCurrentDatabase(spark: SparkSession): Row
+
protected def getGlobalTempViewManager(spark: SparkSession, schemaPattern: String): Seq[String]
// ///////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkCatalogDatabaseOperationSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkCatalogDatabaseOperationSuite.scala
new file mode 100644
index 000000000..46208bff1
--- /dev/null
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkCatalogDatabaseOperationSuite.scala
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.spark.operation
+
+import org.apache.spark.sql.connector.catalog.CatalogPlugin
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+import org.apache.kyuubi.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED
+import org.apache.kyuubi.engine.spark.WithSparkSQLEngine
+import org.apache.kyuubi.engine.spark.shim.SparkCatalogShim
+import org.apache.kyuubi.operation.HiveJDBCTestHelper
+
+class SparkCatalogDatabaseOperationSuite extends WithSparkSQLEngine with HiveJDBCTestHelper {
+
+ override protected def jdbcUrl: String = getJdbcUrl
+
+ override def withKyuubiConf: Map[String, String] =
+ Map(
+ "spark.sql.catalog.dummy" -> classOf[DummyCatalog].getName,
+ ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true")
+
+ test("set/get current catalog") {
+ withJdbcStatement() { statement =>
+ val catalog = statement.getConnection.getCatalog
+ assert(catalog == SparkCatalogShim.SESSION_CATALOG)
+ statement.getConnection.setCatalog("dummy")
+ val changedCatalog = statement.getConnection.getCatalog
+ assert(changedCatalog == "dummy")
+ }
+ }
+
+ test("set/get current database") {
+ withDatabases("test_database") { statement =>
+ statement.execute("create database test_database")
+ val schema = statement.getConnection.getSchema
+ assert(schema == "default")
+ statement.getConnection.setSchema("test_database")
+ val changedSchema = statement.getConnection.getSchema
+ assert(changedSchema == "test_database")
+ }
+ }
+}
+
+class DummyCatalog extends CatalogPlugin {
+ override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {
+ _name = name
+ }
+
+ private var _name: String = null
+
+ override def name(): String = _name
+
+ override def defaultNamespace(): Array[String] = Array("a", "b")
+}
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/TrinoStatement.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/TrinoStatement.scala
index d2b79c09e..cdfec61b9 100644
--- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/TrinoStatement.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/TrinoStatement.scala
@@ -52,6 +52,8 @@ class TrinoStatement(
def getTrinoClient: StatementClient = trino
+ def getCurrentCatalog: String = trinoContext.clientSession.get.getCatalog
+
def getCurrentDatabase: String = trinoContext.clientSession.get.getSchema
def getColumns: List[Column] = {
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentCatalog.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentCatalog.scala
new file mode 100644
index 000000000..2f61f8b13
--- /dev/null
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentCatalog.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.trino.operation
+
+import com.google.common.collect.ImmutableList
+import io.trino.client.{ClientTypeSignature, ClientTypeSignatureParameter, Column}
+import io.trino.client.ClientStandardTypes.VARCHAR
+import io.trino.client.ClientTypeSignature.VARCHAR_UNBOUNDED_LENGTH
+
+import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType}
+import org.apache.kyuubi.session.Session
+
+class GetCurrentCatalog(session: Session)
+ extends TrinoOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val session = trinoContext.clientSession.get
+ val catalog = session.getCatalog
+
+ val clientTypeSignature = new ClientTypeSignature(
+ VARCHAR,
+ ImmutableList.of(ClientTypeSignatureParameter.ofLong(VARCHAR_UNBOUNDED_LENGTH)))
+ val column = new Column("TABLE_CAT", VARCHAR, clientTypeSignature)
+ schema = List[Column](column)
+ iter = new IterableFetchIterator(List(List(catalog)))
+ } catch onError()
+ }
+}
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentDatabase.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentDatabase.scala
new file mode 100644
index 000000000..571fa185d
--- /dev/null
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentDatabase.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.trino.operation
+
+import com.google.common.collect.ImmutableList
+import io.trino.client.{ClientTypeSignature, ClientTypeSignatureParameter, Column}
+import io.trino.client.ClientStandardTypes.VARCHAR
+import io.trino.client.ClientTypeSignature.VARCHAR_UNBOUNDED_LENGTH
+
+import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType}
+import org.apache.kyuubi.session.Session
+
+class GetCurrentDatabase(session: Session)
+ extends TrinoOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val session = trinoContext.clientSession.get
+ val catalog = session.getSchema
+
+ val clientTypeSignature = new ClientTypeSignature(
+ VARCHAR,
+ ImmutableList.of(ClientTypeSignatureParameter.ofLong(VARCHAR_UNBOUNDED_LENGTH)))
+ val column = new Column("TABLE_SCHEM", VARCHAR, clientTypeSignature)
+ schema = List[Column](column)
+ iter = new IterableFetchIterator(List(List(catalog)))
+ } catch onError()
+ }
+}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala
similarity index 56%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala
index 39bef393d..045224d2d 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala
@@ -15,14 +15,23 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.trino.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import io.trino.client.ClientSession
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentCatalog(session: Session, catalog: String)
+ extends TrinoOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val session = trinoContext.clientSession.get
+ var builder = ClientSession.builder(session)
+ builder = builder.withCatalog(catalog)
+ trinoContext.clientSession.set(builder.build())
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala
similarity index 56%
copy from kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
copy to externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala
index 39bef393d..cc9aa1b84 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala
@@ -15,14 +15,23 @@
* limitations under the License.
*/
-package org.apache.kyuubi.config
+package org.apache.kyuubi.engine.trino.operation
-object KyuubiReservedKeys {
- final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user"
- final val KYUUBI_STATEMENT_ID_KEY = "kyuubi.statement.id"
- final val KYUUBI_ENGINE_SUBMIT_TIME_KEY = "kyuubi.engine.submit.time"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_GUID =
- "kyuubi.session.engine.launch.handle.guid"
- final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
- "kyuubi.session.engine.launch.handle.secret"
+import io.trino.client.ClientSession
+
+import org.apache.kyuubi.operation.OperationType
+import org.apache.kyuubi.session.Session
+
+class SetCurrentDatabase(session: Session, database: String)
+ extends TrinoOperation(OperationType.EXECUTE_STATEMENT, session) {
+
+ override protected def runInternal(): Unit = {
+ try {
+ val session = trinoContext.clientSession.get
+ var builder = ClientSession.builder(session)
+ builder = builder.withSchema(database)
+ trinoContext.clientSession.set(builder.build())
+ setHasResultSet(false)
+ } catch onError()
+ }
}
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationManager.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationManager.scala
index 8c0e195f2..63140bd7f 100644
--- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationManager.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationManager.scala
@@ -22,7 +22,7 @@ import java.util
import scala.collection.JavaConverters._
import org.apache.kyuubi.KyuubiSQLException
-import org.apache.kyuubi.config.KyuubiConf.OPERATION_INCREMENTAL_COLLECT
+import org.apache.kyuubi.config.KyuubiConf._
import org.apache.kyuubi.operation.{Operation, OperationManager}
import org.apache.kyuubi.session.Session
@@ -34,12 +34,38 @@ class TrinoOperationManager extends OperationManager("TrinoOperationManager") {
confOverlay: Map[String, String],
runAsync: Boolean,
queryTimeout: Long): Operation = {
+ if (session.sessionManager.getConf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)) {
+ val catalogDatabaseOperation = processCatalogDatabase(session, statement, confOverlay)
+ if (catalogDatabaseOperation != null) {
+ return catalogDatabaseOperation
+ }
+ }
val incrementalCollect = session.sessionManager.getConf.get(OPERATION_INCREMENTAL_COLLECT)
val operation =
new ExecuteStatement(session, statement, runAsync, queryTimeout, incrementalCollect)
addOperation(operation)
}
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ val op = new SetCurrentCatalog(session, catalog)
+ addOperation(op)
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ val op = new GetCurrentCatalog(session)
+ addOperation(op)
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ val op = new SetCurrentDatabase(session, database)
+ addOperation(op)
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ val op = new GetCurrentDatabase(session)
+ addOperation(op)
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val op = new GetTypeInfo(session)
addOperation(op)
diff --git a/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoCatalogDatabaseOperationSuite.scala b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoCatalogDatabaseOperationSuite.scala
new file mode 100644
index 000000000..029fbcf2a
--- /dev/null
+++ b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoCatalogDatabaseOperationSuite.scala
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.trino.operation
+
+import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.engine.trino.WithTrinoEngine
+import org.apache.kyuubi.operation.HiveJDBCTestHelper
+
+class TrinoCatalogDatabaseOperationSuite extends WithTrinoEngine with HiveJDBCTestHelper {
+ override def withKyuubiConf: Map[String, String] = Map(
+ ENGINE_TRINO_CONNECTION_CATALOG.key -> "memory",
+ ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true")
+
+ // use default schema, do not set to 'default', since withSessionHandle strip suffix '/;'
+ override protected val schema = ""
+
+ override protected def jdbcUrl: String = getJdbcUrl
+
+ test("trino - set/get catalog") {
+ withJdbcStatement() { statement =>
+ val catalog = statement.getConnection.getCatalog
+ assert(catalog == "memory")
+ statement.getConnection.setCatalog("system")
+ val changedCatalog = statement.getConnection.getCatalog
+ assert(changedCatalog == "system")
+ }
+ }
+
+ test("trino - set/get database") {
+ withJdbcStatement() { statement =>
+ statement.execute("CREATE SCHEMA IF NOT EXISTS memory.test_trino_db")
+ val schema = statement.getConnection.getSchema
+ assert(schema == "default")
+ statement.getConnection.setSchema("test_trino_db")
+ val changedSchema = statement.getConnection.getSchema
+ assert(changedSchema == "test_trino_db")
+ statement.execute("DROP SCHEMA memory.test_trino_db")
+ }
+ }
+}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
index de9197b80..863734ed3 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
@@ -1571,4 +1571,12 @@ object KyuubiConf {
.version("1.6.0")
.stringConf
.createOptional
+
+ val ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED: ConfigEntry[Boolean] =
+ buildConf("kyuubi.engine.operation.convert.catalog.database.enabled")
+ .doc("When set to true, The engine converts the JDBC methods of set/get Catalog " +
+ "and set/get Schema to the implementation of different engines")
+ .version("1.6.0")
+ .booleanConf
+ .createWithDefault(true)
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
index 39bef393d..190f0b044 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
@@ -25,4 +25,8 @@ object KyuubiReservedKeys {
"kyuubi.session.engine.launch.handle.guid"
final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
"kyuubi.session.engine.launch.handle.secret"
+ final val KYUUBI_OPERATION_SET_CURRENT_CATALOG = "kyuubi.operation.set.current.catalog"
+ final val KYUUBI_OPERATION_GET_CURRENT_CATALOG = "kyuubi.operation.get.current.catalog"
+ final val KYUUBI_OPERATION_SET_CURRENT_DATABASE = "kyuubi.operation.set.current.database"
+ final val KYUUBI_OPERATION_GET_CURRENT_DATABASE = "kyuubi.operation.get.current.database"
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
index e8c107fd0..d5b456234 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
@@ -21,6 +21,7 @@ import org.apache.hive.service.rpc.thrift._
import org.apache.kyuubi.KyuubiSQLException
import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiReservedKeys._
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.operation.OperationState._
import org.apache.kyuubi.operation.log.LogDivertAppender
@@ -51,6 +52,10 @@ abstract class OperationManager(name: String) extends AbstractService(name) {
confOverlay: Map[String, String],
runAsync: Boolean,
queryTimeout: Long): Operation
+ def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation
+ def newGetCurrentCatalogOperation(session: Session): Operation
+ def newSetCurrentDatabaseOperation(session: Session, database: String): Operation
+ def newGetCurrentDatabaseOperation(session: Session): Operation
def newGetTypeInfoOperation(session: Session): Operation
def newGetCatalogsOperation(session: Session): Operation
def newGetSchemasOperation(session: Session, catalog: String, schema: String): Operation
@@ -154,4 +159,30 @@ abstract class OperationManager(name: String) extends AbstractService(name) {
}
}
}
+
+ private val PATTERN_FOR_SET_CATALOG = "_SET_CATALOG"
+ private val PATTERN_FOR_GET_CATALOG = "_GET_CATALOG"
+ private val PATTERN_FOR_SET_SCHEMA = "(?i)use (.*)".r
+ private val PATTERN_FOR_GET_SCHEMA = "select current_database()"
+
+ final def processCatalogDatabase(
+ session: Session,
+ statement: String,
+ confOverlay: Map[String, String]): Operation = {
+ if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_CATALOG)
+ && statement == PATTERN_FOR_SET_CATALOG) {
+ newSetCurrentCatalogOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_CATALOG))
+ } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_CATALOG)
+ && statement == PATTERN_FOR_GET_CATALOG) {
+ newGetCurrentCatalogOperation(session)
+ } else if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_DATABASE)
+ && PATTERN_FOR_SET_SCHEMA.unapplySeq(statement).isDefined) {
+ newSetCurrentDatabaseOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_DATABASE))
+ } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_DATABASE)
+ && PATTERN_FOR_GET_SCHEMA == statement.toLowerCase) {
+ newGetCurrentDatabaseOperation(session)
+ } else {
+ null
+ }
+ }
}
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
index 4f0cfaf13..271e4d2e1 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
@@ -39,6 +39,26 @@ class NoopOperationManager extends OperationManager("noop") {
addOperation(operation)
}
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ val operation = new NoopOperation(OperationType.EXECUTE_STATEMENT, session)
+ addOperation(operation)
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ val operation = new NoopOperation(OperationType.EXECUTE_STATEMENT, session)
+ addOperation(operation)
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ val operation = new NoopOperation(OperationType.EXECUTE_STATEMENT, session)
+ addOperation(operation)
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ val operation = new NoopOperation(OperationType.EXECUTE_STATEMENT, session)
+ addOperation(operation)
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val operation = new NoopOperation(OperationType.GET_TYPE_INFO, session)
addOperation(operation)
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 c3cd2962e..5aa94e947 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
@@ -1094,6 +1094,10 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
return new KyuubiStatement(this, client, sessHandle, fetchSize);
}
+ private KyuubiStatement createKyuubiStatement() throws SQLException {
+ return ((KyuubiStatement) createStatement());
+ }
+
/*
* (non-Javadoc)
*
@@ -1164,7 +1168,18 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
@Override
public String getCatalog() throws SQLException {
- return "";
+ if (isClosed) {
+ throw new SQLException("Connection is closed");
+ }
+ try (KyuubiStatement stmt = createKyuubiStatement();
+ ResultSet res = stmt.executeGetCurrentCatalog("_GET_CATALOG")) {
+ if (!res.next()) {
+ throw new SQLException("Failed to get catalog information");
+ }
+ return res.getString(1);
+ } catch (Exception e) {
+ return "";
+ }
}
/*
@@ -1225,8 +1240,8 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
if (isClosed) {
throw new SQLException("Connection is closed");
}
- try (Statement stmt = createStatement();
- ResultSet res = stmt.executeQuery("SELECT current_database()")) {
+ try (KyuubiStatement stmt = createKyuubiStatement();
+ ResultSet res = stmt.executeGetCurrentDatabase("SELECT current_database()")) {
if (!res.next()) {
throw new SQLException("Failed to get schema information");
}
@@ -1514,12 +1529,14 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
@Override
public void setCatalog(String catalog) throws SQLException {
- // Per JDBC spec, if the driver does not support catalogs,
- // it will silently ignore this request.
if (isClosed) {
throw new SQLException("Connection is closed");
}
- return;
+ try (KyuubiStatement stmt = createKyuubiStatement()) {
+ stmt.executeSetCurrentCatalog("_SET_CATALOG", catalog);
+ } catch (SQLException e) {
+
+ }
}
/*
@@ -1641,8 +1658,8 @@ public class KyuubiConnection implements java.sql.Connection, KyuubiLoggable {
if (schema == null || schema.isEmpty()) {
throw new SQLException("Schema name is null or empty");
}
- try (Statement stmt = createStatement()) {
- stmt.execute("use " + schema);
+ try (KyuubiStatement stmt = createKyuubiStatement()) {
+ stmt.executeSetCurrentDatabase("use " + schema, schema);
}
}
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 31156f5fb..4f13a187a 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
@@ -512,6 +512,36 @@ public class KyuubiStatement implements java.sql.Statement, KyuubiLoggable {
return resultSet;
}
+ public void executeSetCurrentCatalog(String sql, String catalog) throws SQLException {
+ if (executeWithConfOverlay(
+ sql, Collections.singletonMap("kyuubi.operation.set.current.catalog", catalog))) {
+ resultSet.close();
+ }
+ }
+
+ public ResultSet executeGetCurrentCatalog(String sql) throws SQLException {
+ if (!executeWithConfOverlay(
+ sql, Collections.singletonMap("kyuubi.operation.get.current.catalog", ""))) {
+ throw new SQLException("The query did not generate a result set!");
+ }
+ return resultSet;
+ }
+
+ public void executeSetCurrentDatabase(String sql, String database) throws SQLException {
+ if (executeWithConfOverlay(
+ sql, Collections.singletonMap("kyuubi.operation.set.current.database", database))) {
+ resultSet.close();
+ }
+ }
+
+ public ResultSet executeGetCurrentDatabase(String sql) throws SQLException {
+ if (!executeWithConfOverlay(
+ sql, Collections.singletonMap("kyuubi.operation.get.current.database", ""))) {
+ throw new SQLException("The query did not generate a result set!");
+ }
+ return resultSet;
+ }
+
/*
* (non-Javadoc)
*
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
index d50b6afb3..7ad1c99e9 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
@@ -71,6 +71,23 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam
operation
}
+ // The server does not use these 4 operations
+ override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = {
+ null
+ }
+
+ override def newGetCurrentCatalogOperation(session: Session): Operation = {
+ null
+ }
+
+ override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = {
+ null
+ }
+
+ override def newGetCurrentDatabaseOperation(session: Session): Operation = {
+ null
+ }
+
override def newGetTypeInfoOperation(session: Session): Operation = {
val operation = new GetTypeInfo(session)
addOperation(operation)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationManagerSuite.scala
index e7082463b..1bd750e6a 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationManagerSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationManagerSuite.scala
@@ -21,7 +21,7 @@ import java.sql.SQLTimeoutException
import org.apache.kyuubi.WithKyuubiServer
import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.config.KyuubiConf.OPERATION_QUERY_TIMEOUT
+import org.apache.kyuubi.config.KyuubiConf._
class KyuubiOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelper {
override protected val conf: KyuubiConf = {
@@ -60,5 +60,35 @@ class KyuubiOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelp
}
}
+ test("set/get current catalog") {
+ withSessionConf()(
+ Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true"))(
+ Map.empty) {
+ withJdbcStatement() { statement =>
+ val catalog = statement.getConnection.getCatalog
+ assert(catalog == "spark_catalog")
+ // The server starts the spark engine without other catalogs
+ statement.getConnection.setCatalog("dummy_catalog")
+ val changedCatalog = statement.getConnection.getCatalog
+ assert(changedCatalog == "spark_catalog")
+ }
+ }
+ }
+
+ test("set/get current database") {
+ withSessionConf()(
+ Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true"))(
+ Map.empty) {
+ withDatabases("test_database") { statement =>
+ statement.execute("create database test_database")
+ val schema = statement.getConnection.getSchema
+ assert(schema == "default")
+ statement.getConnection.setSchema("test_database")
+ val changedSchema = statement.getConnection.getSchema
+ assert(changedSchema == "test_database")
+ }
+ }
+ }
+
override protected def jdbcUrl: String = getJdbcUrl
}