From 21e578701ba860f4df0e3086649319ff8f7c7b7a Mon Sep 17 00:00:00 2001 From: sychen Date: Tue, 24 May 2022 05:13:18 +0800 Subject: [PATCH 01/13] set/get catalog/database operators --- .../operation/FlinkSQLOperationManager.scala | 24 ++++++++++ .../flink/operation/GetCurrentCatalog.scala | 35 ++++++++++++++ .../flink/operation/GetCurrentDatabase.scala | 35 ++++++++++++++ .../flink/operation/SetCurrentCatalog.scala | 35 ++++++++++++++ .../flink/operation/SetCurrentDatabase.scala | 35 ++++++++++++++ .../hive/operation/GetCurrentCatalog.scala | 30 ++++++++++++ .../hive/operation/GetCurrentDatabase.scala | 37 +++++++++++++++ .../hive/operation/HiveOperationManager.scala | 24 ++++++++++ .../hive/operation/SetCurrentCatalog.scala | 30 ++++++++++++ .../hive/operation/SetCurrentDatabase.scala | 37 +++++++++++++++ .../spark/operation/GetCurrentCatalog.scala | 40 ++++++++++++++++ .../spark/operation/GetCurrentDatabase.scala | 40 ++++++++++++++++ .../spark/operation/SetCurrentCatalog.scala | 41 +++++++++++++++++ .../spark/operation/SetCurrentDatabase.scala | 41 +++++++++++++++++ .../operation/SparkSQLOperationManager.scala | 24 ++++++++++ .../engine/spark/shim/CatalogShim_v2_4.scala | 14 ++++++ .../engine/spark/shim/CatalogShim_v3_0.scala | 16 +++++++ .../engine/spark/shim/SparkCatalogShim.scala | 8 ++++ .../kyuubi/engine/trino/TrinoStatement.scala | 2 + .../trino/operation/GetCurrentCatalog.scala | 44 ++++++++++++++++++ .../trino/operation/GetCurrentDatabase.scala | 44 ++++++++++++++++++ .../trino/operation/SetCurrentCatalog.scala | 46 +++++++++++++++++++ .../trino/operation/SetCurrentDatabase.scala | 46 +++++++++++++++++++ .../operation/TrinoOperationManager.scala | 24 ++++++++++ .../kyuubi/config/KyuubiReservedKeys.scala | 4 ++ .../kyuubi/operation/OperationManager.scala | 20 ++++++++ .../operation/NoopOperationManager.scala | 20 ++++++++ .../kyuubi/jdbc/hive/KyuubiConnection.java | 35 ++++++++++---- .../kyuubi/jdbc/hive/KyuubiStatement.java | 32 +++++++++++++ .../kyuubi/operation/GetCurrentCatalog.scala | 32 +++++++++++++ .../kyuubi/operation/GetCurrentDatabase.scala | 32 +++++++++++++ .../operation/KyuubiOperationManager.scala | 24 ++++++++++ .../kyuubi/operation/SetCurrentCatalog.scala | 32 +++++++++++++ .../kyuubi/operation/SetCurrentDatabase.scala | 32 +++++++++++++ 34 files changed, 1007 insertions(+), 8 deletions(-) create mode 100644 externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala create mode 100644 externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala create mode 100644 externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala create mode 100644 externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala create mode 100644 externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala create mode 100644 externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala create mode 100644 externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala create mode 100644 externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala create mode 100644 externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala create mode 100644 externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala create mode 100644 externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala create mode 100644 externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala create mode 100644 externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentCatalog.scala create mode 100644 externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/GetCurrentDatabase.scala create mode 100644 externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala create mode 100644 externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala create mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala create mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala create mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala create mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala 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 9a98fe1b78f..fcd60ac132c 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 @@ -42,6 +42,10 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { + val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + if (catalogDatabaseOperation != null) { + return catalogDatabaseOperation + } val flinkSession = session.asInstanceOf[FlinkSessionImpl] val mode = flinkSession.sessionContext.getConfigMap.getOrDefault( OPERATION_PLAN_ONLY_MODE.key, @@ -59,6 +63,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/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala new file mode 100644 index 00000000000..2ddb063615c --- /dev/null +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentCatalog.scala @@ -0,0 +1,35 @@ +/* + * 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.flink.operation + +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/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala new file mode 100644 index 00000000000..1e392511ef8 --- /dev/null +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/GetCurrentDatabase.scala @@ -0,0 +1,35 @@ +/* + * 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.flink.operation + +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/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala new file mode 100644 index 00000000000..f79b9bbedb1 --- /dev/null +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala @@ -0,0 +1,35 @@ +/* + * 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.flink.operation + +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 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) + resultSet = ResultSetUtil.stringListToResultSet(List(""), TABLE_CAT) + } catch onError() + } +} diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala new file mode 100644 index 00000000000..c7bb24df32e --- /dev/null +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala @@ -0,0 +1,35 @@ +/* + * 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.flink.operation + +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 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) + resultSet = ResultSetUtil.stringListToResultSet(List(""), TABLE_SCHEM) + } catch onError() + } +} diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala new file mode 100644 index 00000000000..c334eac5563 --- /dev/null +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala @@ -0,0 +1,30 @@ +/* + * 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.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.newGetCatalogsOperation(hive) +} diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala new file mode 100644 index 00000000000..3bdbffa1531 --- /dev/null +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentDatabase.scala @@ -0,0 +1,37 @@ +/* + * 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 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 cdfea54fa41..33026ff7ef6 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 @@ -36,10 +36,34 @@ class HiveOperationManager() extends OperationManager("HiveOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { + val catalogDatabaseOperation = processCatalogDatabase(session, 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/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala new file mode 100644 index 00000000000..f4881d3fff1 --- /dev/null +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala @@ -0,0 +1,30 @@ +/* + * 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.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.newGetCatalogsOperation(hive) +} diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala new file mode 100644 index 00000000000..c23127d384a --- /dev/null +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentDatabase.scala @@ -0,0 +1,37 @@ +/* + * 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 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-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala new file mode 100644 index 00000000000..47b386016d1 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentCatalog.scala @@ -0,0 +1,40 @@ +/* + * 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.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/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala new file mode 100644 index 00000000000..954e23898ae --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/GetCurrentDatabase.scala @@ -0,0 +1,40 @@ +/* + * 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.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/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala new file mode 100644 index 00000000000..1bd57971a73 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala @@ -0,0 +1,41 @@ +/* + * 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.Row +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 SetCurrentCatalog(session: Session, catalog: String) + extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def resultSchema: StructType = { + new StructType().add(TABLE_CAT, "string") + } + + override protected def runInternal(): Unit = { + try { + SparkCatalogShim().setCurrentCatalog(spark, catalog) + iter = new IterableFetchIterator(Seq(Row(""))) + } catch onError() + } +} diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala new file mode 100644 index 00000000000..52a48ea8a11 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala @@ -0,0 +1,41 @@ +/* + * 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.Row +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 SetCurrentDatabase(session: Session, database: String) + extends SparkOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def resultSchema: StructType = { + new StructType().add(TABLE_SCHEM, "string") + } + + override protected def runInternal(): Unit = { + try { + SparkCatalogShim().setCurrentDatabase(spark, database) + iter = new IterableFetchIterator(Seq(Row(""))) + } 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 bc7b84390aa..e35ff2e837a 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 @@ -53,6 +53,10 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n runAsync: Boolean, queryTimeout: Long): Operation = { val spark = session.asInstanceOf[SparkSessionImpl].spark + val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + if (catalogDatabaseOperation != null) { + return catalogDatabaseOperation + } val lang = confOverlay.getOrElse( OPERATION_LANGUAGE.key, spark.conf.get(OPERATION_LANGUAGE.key, operationLanguageDefault)) @@ -75,6 +79,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 cc5a25b0c29..5977cd415b0 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 db0335849c4..4dd3a230b95 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,14 @@ class CatalogShim_v3_0 extends CatalogShim_v2_4 { spark.sessionState.catalogManager.isCatalogRegistered(catalog) } + override def setCurrentCatalog(spark: SparkSession, catalog: String): Unit = { + spark.sessionState.catalogManager.setCurrentCatalog(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 +130,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 831b39c7457..bc5792823f7 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-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 acd1f0408ab..f35b0348b04 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 @@ -58,6 +58,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 00000000000..2f61f8b131d --- /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 00000000000..571fa185d37 --- /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/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala new file mode 100644 index 00000000000..a350d99dda5 --- /dev/null +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala @@ -0,0 +1,46 @@ +/* + * 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.{ClientSession, 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 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()) + + 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(""))) + } catch onError() + } +} diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala new file mode 100644 index 00000000000..1a8f049ad55 --- /dev/null +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala @@ -0,0 +1,46 @@ +/* + * 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.{ClientSession, 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 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()) + + 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(""))) + } 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 8c0e195f2a5..811cda470e3 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 @@ -34,12 +34,36 @@ class TrinoOperationManager extends OperationManager("TrinoOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { + val catalogDatabaseOperation = processCatalogDatabase(session, 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/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala index 39bef393d62..190f0b044b8 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 e8c107fd0ef..b5b150ece82 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.{KYUUBI_OPERATION_GET_CURRENT_CATALOG, KYUUBI_OPERATION_GET_CURRENT_DATABASE, KYUUBI_OPERATION_SET_CURRENT_CATALOG, KYUUBI_OPERATION_SET_CURRENT_DATABASE} 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,19 @@ abstract class OperationManager(name: String) extends AbstractService(name) { } } } + final def processCatalogDatabase( + session: Session, + confOverlay: Map[String, String]): Operation = { + if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_CATALOG)) { + newSetCurrentCatalogOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_CATALOG)) + } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_CATALOG)) { + newGetCurrentCatalogOperation(session) + } else if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_DATABASE)) { + newSetCurrentDatabaseOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_DATABASE)) + } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_DATABASE)) { + 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 4f0cfaf1373..271e4d2e1d1 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 c3cd2962edb..2b2f5a2c607 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 Statement createStatement() throws SQLException { return new KyuubiStatement(this, client, sessHandle, fetchSize); } + private KyuubiStatement createKyuubiStatement() throws SQLException { + return ((KyuubiStatement) createStatement()); + } + /* * (non-Javadoc) * @@ -1164,7 +1168,16 @@ public boolean getAutoCommit() throws SQLException { @Override public String getCatalog() throws SQLException { - return ""; + if (isClosed) { + throw new SQLException("Connection is closed"); + } + try (KyuubiStatement stmt = createKyuubiStatement(); + ResultSet res = stmt.executeGetCurrentCatalog()) { + if (!res.next()) { + throw new SQLException("Failed to get catalog information"); + } + return res.getString(1); + } } /* @@ -1225,8 +1238,8 @@ public String getSchema() throws SQLException { 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()) { if (!res.next()) { throw new SQLException("Failed to get schema information"); } @@ -1514,12 +1527,15 @@ public void setAutoCommit(boolean autoCommit) throws SQLException { @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(); + ResultSet res = stmt.executeSetCurrentCatalog(catalog)) { + if (!res.next()) { + throw new SQLException("Failed to set catalog"); + } + } } /* @@ -1641,8 +1657,11 @@ public void setSchema(String schema) throws SQLException { 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(); + ResultSet res = stmt.executeSetCurrentDatabase(schema)) { + if (!res.next()) { + throw new SQLException("Failed to set 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 31156f5fb28..ae0b34b0b07 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,38 @@ public ResultSet executeScala(String code) throws SQLException { return resultSet; } + public ResultSet executeSetCurrentCatalog(String catalog) throws SQLException { + if (!executeWithConfOverlay( + "", Collections.singletonMap("kyuubi.operation.set.current.catalog", catalog))) { + throw new SQLException("The query did not generate a result set!"); + } + return resultSet; + } + + public ResultSet executeGetCurrentCatalog() throws SQLException { + if (!executeWithConfOverlay( + "", Collections.singletonMap("kyuubi.operation.get.current.catalog", ""))) { + throw new SQLException("The query did not generate a result set!"); + } + return resultSet; + } + + public ResultSet executeSetCurrentDatabase(String database) throws SQLException { + if (!executeWithConfOverlay( + "", Collections.singletonMap("kyuubi.operation.set.current.database", database))) { + throw new SQLException("The query did not generate a result set!"); + } + return resultSet; + } + + public ResultSet executeGetCurrentDatabase() throws SQLException { + if (!executeWithConfOverlay( + "", 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/GetCurrentCatalog.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala new file mode 100644 index 00000000000..05045bbc984 --- /dev/null +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala @@ -0,0 +1,32 @@ +/* + * 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.operation + +import org.apache.kyuubi.config.KyuubiReservedKeys._ +import org.apache.kyuubi.session.Session + +class GetCurrentCatalog(session: Session) + extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def runInternal(): Unit = { + try { + val conf = Map(KYUUBI_OPERATION_GET_CURRENT_CATALOG -> "") + _remoteOpHandle = client.executeStatement("", conf, false, 0L) + } catch onError() + } +} diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala new file mode 100644 index 00000000000..88d2bae7c32 --- /dev/null +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala @@ -0,0 +1,32 @@ +/* + * 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.operation + +import org.apache.kyuubi.config.KyuubiReservedKeys._ +import org.apache.kyuubi.session.Session + +class GetCurrentDatabase(session: Session) + extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def runInternal(): Unit = { + try { + val conf = Map(KYUUBI_OPERATION_GET_CURRENT_DATABASE -> "") + _remoteOpHandle = client.executeStatement("", conf, false, 0L) + } catch onError() + } +} 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 d50b6afb369..78e470ef4de 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 @@ -58,6 +58,10 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { + val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + if (catalogDatabaseOperation != null) { + return catalogDatabaseOperation + } val operation = new ExecuteStatement(session, statement, confOverlay, runAsync, getQueryTimeout(queryTimeout)) addOperation(operation) @@ -71,6 +75,26 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam operation } + override def newSetCurrentCatalogOperation(session: Session, catalog: String): Operation = { + val operation = new SetCurrentCatalog(session, catalog) + addOperation(operation) + } + + override def newGetCurrentCatalogOperation(session: Session): Operation = { + val operation = new GetCurrentCatalog(session) + addOperation(operation) + } + + override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = { + val operation = new SetCurrentDatabase(session, database) + addOperation(operation) + } + + override def newGetCurrentDatabaseOperation(session: Session): Operation = { + val operation = new GetCurrentDatabase(session) + addOperation(operation) + } + override def newGetTypeInfoOperation(session: Session): Operation = { val operation = new GetTypeInfo(session) addOperation(operation) diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala new file mode 100644 index 00000000000..a96a93e52bc --- /dev/null +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala @@ -0,0 +1,32 @@ +/* + * 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.operation + +import org.apache.kyuubi.config.KyuubiReservedKeys._ +import org.apache.kyuubi.session.Session + +class SetCurrentCatalog(session: Session, catalog: String) + extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def runInternal(): Unit = { + try { + val conf = Map(KYUUBI_OPERATION_SET_CURRENT_CATALOG -> catalog) + _remoteOpHandle = client.executeStatement("", conf, false, 0L) + } catch onError() + } +} diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala new file mode 100644 index 00000000000..06636d45647 --- /dev/null +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala @@ -0,0 +1,32 @@ +/* + * 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.operation + +import org.apache.kyuubi.config.KyuubiReservedKeys._ +import org.apache.kyuubi.session.Session + +class SetCurrentDatabase(session: Session, database: String) + extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { + + override protected def runInternal(): Unit = { + try { + val conf = Map(KYUUBI_OPERATION_SET_CURRENT_DATABASE -> database) + _remoteOpHandle = client.executeStatement("", conf, false, 0L) + } catch onError() + } +} From 9877038f8ebe5c7ec43fb71bfc44395c9af10c6c Mon Sep 17 00:00:00 2001 From: sychen Date: Tue, 24 May 2022 21:50:32 +0800 Subject: [PATCH 02/13] add UT --- .../flink/operation/FlinkOperationSuite.scala | 24 +++++++ .../hive/operation/GetCurrentCatalog.scala | 9 ++- .../hive/operation/SetCurrentCatalog.scala | 9 ++- .../SparkCatalogDatabaseOperationSuite.scala | 67 +++++++++++++++++++ .../trino/operation/TrinoOperationSuite.scala | 22 ++++++ .../org/apache/kyuubi/HiveEngineTests.scala | 22 ++++++ .../kyuubi/jdbc/hive/KyuubiConnection.java | 14 ++-- .../kyuubi/jdbc/hive/KyuubiStatement.java | 14 ++-- .../kyuubi/operation/SetCurrentCatalog.scala | 1 + .../kyuubi/operation/SetCurrentDatabase.scala | 1 + .../KyuubiOperationManagerSuite.scala | 23 +++++++ 11 files changed, 186 insertions(+), 20 deletions(-) create mode 100644 externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkCatalogDatabaseOperationSuite.scala 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 41998238712..ab65cafd1fb 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 @@ -665,6 +665,18 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { }) } + test("execute statement - set/get catalog") { + 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 +686,18 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { }) } + test("execute statement - set/get database") { + 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 => diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala index c334eac5563..ee6d44df053 100644 --- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/GetCurrentCatalog.scala @@ -17,6 +17,8 @@ package org.apache.kyuubi.engine.hive.operation +import scala.collection.JavaConverters._ + import org.apache.hive.service.cli.operation.Operation import org.apache.kyuubi.operation.OperationType @@ -26,5 +28,10 @@ class GetCurrentCatalog(session: Session) extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) { // Hive does not support catalog override val internalHiveOperation: Operation = - delegatedOperationManager.newGetCatalogsOperation(hive) + delegatedOperationManager.newExecuteStatementOperation( + hive, + "SELECT '' AS TABLE_CAT", + Map.empty[String, String].asJava, + false, + 0) } diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala index f4881d3fff1..c6df59a8941 100644 --- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala @@ -17,6 +17,8 @@ package org.apache.kyuubi.engine.hive.operation +import scala.collection.JavaConverters._ + import org.apache.hive.service.cli.operation.Operation import org.apache.kyuubi.operation.OperationType @@ -26,5 +28,10 @@ class SetCurrentCatalog(session: Session, catalog: String) extends HiveOperation(OperationType.EXECUTE_STATEMENT, session) { // Hive does not support catalog override val internalHiveOperation: Operation = - delegatedOperationManager.newGetCatalogsOperation(hive) + delegatedOperationManager.newExecuteStatementOperation( + hive, + "SELECT '' AS TABLE_CAT", + Map.empty[String, String].asJava, + false, + 0) } 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 00000000000..97380b06bd9 --- /dev/null +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/SparkCatalogDatabaseOperationSuite.scala @@ -0,0 +1,67 @@ +/* + * 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.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) + + 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/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala index d3680c30aaa..f5d63fdcd21 100644 --- a/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala +++ b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala @@ -738,4 +738,26 @@ class TrinoOperationSuite extends WithTrinoEngine with TrinoQueryTests { assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS) } } + + 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/test/scala/org/apache/kyuubi/HiveEngineTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala index 9eb4a24405b..a6049ae2cb7 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala @@ -452,4 +452,26 @@ trait HiveEngineTests extends HiveJDBCTestHelper { assert(resultSet.getString(1) === "hive.query.name=test") } } + + test("test set/get catalog") { + 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") { + 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/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 2b2f5a2c607..b34957354ec 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 @@ -1530,11 +1530,8 @@ public void setCatalog(String catalog) throws SQLException { if (isClosed) { throw new SQLException("Connection is closed"); } - try (KyuubiStatement stmt = createKyuubiStatement(); - ResultSet res = stmt.executeSetCurrentCatalog(catalog)) { - if (!res.next()) { - throw new SQLException("Failed to set catalog"); - } + try (KyuubiStatement stmt = createKyuubiStatement()) { + stmt.executeSetCurrentCatalog(catalog); } } @@ -1657,11 +1654,8 @@ public void setSchema(String schema) throws SQLException { if (schema == null || schema.isEmpty()) { throw new SQLException("Schema name is null or empty"); } - try (KyuubiStatement stmt = createKyuubiStatement(); - ResultSet res = stmt.executeSetCurrentDatabase(schema)) { - if (!res.next()) { - throw new SQLException("Failed to set schema"); - } + try (KyuubiStatement stmt = createKyuubiStatement()) { + stmt.executeSetCurrentDatabase(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 ae0b34b0b07..2595c029eb5 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,12 +512,11 @@ public ResultSet executeScala(String code) throws SQLException { return resultSet; } - public ResultSet executeSetCurrentCatalog(String catalog) throws SQLException { - if (!executeWithConfOverlay( + public void executeSetCurrentCatalog(String catalog) throws SQLException { + if (executeWithConfOverlay( "", Collections.singletonMap("kyuubi.operation.set.current.catalog", catalog))) { - throw new SQLException("The query did not generate a result set!"); + resultSet.close(); } - return resultSet; } public ResultSet executeGetCurrentCatalog() throws SQLException { @@ -528,12 +527,11 @@ public ResultSet executeGetCurrentCatalog() throws SQLException { return resultSet; } - public ResultSet executeSetCurrentDatabase(String database) throws SQLException { - if (!executeWithConfOverlay( + public void executeSetCurrentDatabase(String database) throws SQLException { + if (executeWithConfOverlay( "", Collections.singletonMap("kyuubi.operation.set.current.database", database))) { - throw new SQLException("The query did not generate a result set!"); + resultSet.close(); } - return resultSet; } public ResultSet executeGetCurrentDatabase() throws SQLException { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala index a96a93e52bc..6a7c36eee5a 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala @@ -27,6 +27,7 @@ class SetCurrentCatalog(session: Session, catalog: String) try { val conf = Map(KYUUBI_OPERATION_SET_CURRENT_CATALOG -> catalog) _remoteOpHandle = client.executeStatement("", conf, false, 0L) + setHasResultSet(_remoteOpHandle.isHasResultSet) } catch onError() } } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala index 06636d45647..0da3f1b561b 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala @@ -27,6 +27,7 @@ class SetCurrentDatabase(session: Session, database: String) try { val conf = Map(KYUUBI_OPERATION_SET_CURRENT_DATABASE -> database) _remoteOpHandle = client.executeStatement("", conf, false, 0L) + setHasResultSet(_remoteOpHandle.isHasResultSet) } catch onError() } } 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 e7082463b98..0ac2e6c3c66 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 @@ -17,6 +17,7 @@ package org.apache.kyuubi.operation +import java.sql.SQLException import java.sql.SQLTimeoutException import org.apache.kyuubi.WithKyuubiServer @@ -60,5 +61,27 @@ class KyuubiOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelp } } + test("set/get current catalog") { + withJdbcStatement() { statement => + val catalog = statement.getConnection.getCatalog + assert(catalog == "spark_catalog") + statement.getConnection.setCatalog("dummy_catalog") + // The server starts the spark engine without other catalogs + val e = intercept[SQLException](statement.getConnection.getCatalog) + assert(e.getMessage.contains("dummy_catalog")) + } + } + + 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") + } + } + override protected def jdbcUrl: String = getJdbcUrl } From 4982ec71b93c6182148f78599239d34695374ef7 Mon Sep 17 00:00:00 2001 From: sychen Date: Tue, 24 May 2022 22:07:03 +0800 Subject: [PATCH 03/13] hive engine 1.8 --- .../src/test/scala/org/apache/kyuubi/HiveEngineTests.scala | 2 ++ 1 file changed, 2 insertions(+) diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala index a6049ae2cb7..b957a8ee924 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala @@ -454,6 +454,7 @@ trait HiveEngineTests extends HiveJDBCTestHelper { } test("test set/get catalog") { + assume(SystemUtils.isJavaVersionAtMost(JavaVersion.JAVA_1_8)) withJdbcStatement()({ statement => val catalog = statement.getConnection.getCatalog assert(catalog == "") @@ -464,6 +465,7 @@ trait HiveEngineTests extends HiveJDBCTestHelper { } 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 From 7f13ad8d576ddf57c55d12aed6c45f80be9e2a7b Mon Sep 17 00:00:00 2001 From: sychen Date: Tue, 24 May 2022 23:18:59 +0800 Subject: [PATCH 04/13] fix server spark catalog UT --- .../kyuubi/operation/KyuubiOperationManagerSuite.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 0ac2e6c3c66..abae65d0d55 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 @@ -65,9 +65,11 @@ class KyuubiOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelp withJdbcStatement() { statement => val catalog = statement.getConnection.getCatalog assert(catalog == "spark_catalog") - statement.getConnection.setCatalog("dummy_catalog") // The server starts the spark engine without other catalogs - val e = intercept[SQLException](statement.getConnection.getCatalog) + val e = intercept[SQLException] { + statement.getConnection.setCatalog("dummy_catalog") + statement.getConnection.getCatalog + } assert(e.getMessage.contains("dummy_catalog")) } } From 45caf5522a4ed1849c78cbbbffed5e2f38b3bf06 Mon Sep 17 00:00:00 2001 From: sychen Date: Wed, 25 May 2022 12:33:40 +0800 Subject: [PATCH 05/13] fix import --- .../scala/org/apache/kyuubi/operation/OperationManager.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b5b150ece82..b533167b493 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,7 +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.{KYUUBI_OPERATION_GET_CURRENT_CATALOG, KYUUBI_OPERATION_GET_CURRENT_DATABASE, KYUUBI_OPERATION_SET_CURRENT_CATALOG, KYUUBI_OPERATION_SET_CURRENT_DATABASE} +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 From 698fbbdd2a7c49793b8da7a1cdcc104167e02da4 Mon Sep 17 00:00:00 2001 From: sychen Date: Wed, 25 May 2022 15:03:24 +0800 Subject: [PATCH 06/13] set catalog/schema empty resultset --- .../flink/operation/SetCurrentCatalog.scala | 4 +--- .../flink/operation/SetCurrentDatabase.scala | 4 +--- .../engine/hive/operation/SetCurrentCatalog.scala | 7 ++++++- .../spark/operation/SetCurrentCatalog.scala | 8 +++----- .../spark/operation/SetCurrentDatabase.scala | 8 +++----- .../trino/operation/SetCurrentCatalog.scala | 15 +++------------ .../trino/operation/SetCurrentDatabase.scala | 15 +++------------ 7 files changed, 20 insertions(+), 41 deletions(-) diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala index f79b9bbedb1..6090d504044 100644 --- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentCatalog.scala @@ -17,9 +17,7 @@ package org.apache.kyuubi.engine.flink.operation -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 SetCurrentCatalog(session: Session, catalog: String) @@ -29,7 +27,7 @@ class SetCurrentCatalog(session: Session, catalog: String) try { val tableEnv = sessionContext.getExecutionContext.getTableEnvironment tableEnv.useCatalog(catalog) - resultSet = ResultSetUtil.stringListToResultSet(List(""), TABLE_CAT) + setHasResultSet(false) } catch onError() } } diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala index c7bb24df32e..633a1197478 100644 --- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/SetCurrentDatabase.scala @@ -17,9 +17,7 @@ package org.apache.kyuubi.engine.flink.operation -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 SetCurrentDatabase(session: Session, database: String) @@ -29,7 +27,7 @@ class SetCurrentDatabase(session: Session, database: String) try { val tableEnv = sessionContext.getExecutionContext.getTableEnvironment tableEnv.useDatabase(database) - resultSet = ResultSetUtil.stringListToResultSet(List(""), TABLE_SCHEM) + setHasResultSet(false) } catch onError() } } diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala index c6df59a8941..24340fb2c7a 100644 --- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/SetCurrentCatalog.scala @@ -30,8 +30,13 @@ class SetCurrentCatalog(session: Session, catalog: String) override val internalHiveOperation: Operation = delegatedOperationManager.newExecuteStatementOperation( hive, - "SELECT '' AS TABLE_CAT", + "SELECT ''", Map.empty[String, String].asJava, false, 0) + + override def runInternal(): Unit = { + super.runInternal() + setHasResultSet(false) + } } diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala index 1bd57971a73..dd0bf987c8c 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentCatalog.scala @@ -17,25 +17,23 @@ package org.apache.kyuubi.engine.spark.operation -import org.apache.spark.sql.Row 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.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().add(TABLE_CAT, "string") + new StructType() } override protected def runInternal(): Unit = { try { SparkCatalogShim().setCurrentCatalog(spark, catalog) - iter = new IterableFetchIterator(Seq(Row(""))) + setHasResultSet(false) } catch onError() } } diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala index 52a48ea8a11..b4efc15b091 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SetCurrentDatabase.scala @@ -17,25 +17,23 @@ package org.apache.kyuubi.engine.spark.operation -import org.apache.spark.sql.Row 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.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().add(TABLE_SCHEM, "string") + new StructType() } override protected def runInternal(): Unit = { try { SparkCatalogShim().setCurrentDatabase(spark, database) - iter = new IterableFetchIterator(Seq(Row(""))) + setHasResultSet(false) } catch onError() } } diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala index a350d99dda5..045224d2d1a 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala @@ -17,12 +17,9 @@ package org.apache.kyuubi.engine.trino.operation -import com.google.common.collect.ImmutableList -import io.trino.client.{ClientSession, ClientTypeSignature, ClientTypeSignatureParameter, Column} -import io.trino.client.ClientStandardTypes.VARCHAR -import io.trino.client.ClientTypeSignature.VARCHAR_UNBOUNDED_LENGTH +import io.trino.client.ClientSession -import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType} +import org.apache.kyuubi.operation.OperationType import org.apache.kyuubi.session.Session class SetCurrentCatalog(session: Session, catalog: String) @@ -34,13 +31,7 @@ class SetCurrentCatalog(session: Session, catalog: String) var builder = ClientSession.builder(session) builder = builder.withCatalog(catalog) trinoContext.clientSession.set(builder.build()) - - 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(""))) + setHasResultSet(false) } catch onError() } } diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala index 1a8f049ad55..cc9aa1b84bd 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala @@ -17,12 +17,9 @@ package org.apache.kyuubi.engine.trino.operation -import com.google.common.collect.ImmutableList -import io.trino.client.{ClientSession, ClientTypeSignature, ClientTypeSignatureParameter, Column} -import io.trino.client.ClientStandardTypes.VARCHAR -import io.trino.client.ClientTypeSignature.VARCHAR_UNBOUNDED_LENGTH +import io.trino.client.ClientSession -import org.apache.kyuubi.operation.{IterableFetchIterator, OperationType} +import org.apache.kyuubi.operation.OperationType import org.apache.kyuubi.session.Session class SetCurrentDatabase(session: Session, database: String) @@ -34,13 +31,7 @@ class SetCurrentDatabase(session: Session, database: String) var builder = ClientSession.builder(session) builder = builder.withSchema(database) trinoContext.clientSession.set(builder.build()) - - 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(""))) + setHasResultSet(false) } catch onError() } } From f62461ce7ee42090e0f799ec6fc7a1f89193333a Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 01:55:10 +0800 Subject: [PATCH 07/13] use placeholders --- docs/deployment/settings.md | 1 + .../operation/FlinkSQLOperationManager.scala | 2 +- .../FlinkCatalogDatabaseOperationSuite.scala | 54 +++++++++++++++ .../flink/operation/FlinkOperationSuite.scala | 24 ------- .../hive/operation/HiveOperationManager.scala | 2 +- .../HiveCatalogDatabaseOperationSuite.scala | 68 +++++++++++++++++++ .../operation/SparkSQLOperationManager.scala | 2 +- .../SparkCatalogDatabaseOperationSuite.scala | 4 +- .../operation/TrinoOperationManager.scala | 2 +- .../TrinoCatalogDatabaseOperationSuite.scala | 55 +++++++++++++++ .../trino/operation/TrinoOperationSuite.scala | 22 ------ .../org/apache/kyuubi/config/KyuubiConf.scala | 8 +++ .../kyuubi/config/KyuubiReservedKeys.scala | 4 -- .../kyuubi/operation/OperationManager.scala | 32 ++++++--- .../org/apache/kyuubi/HiveEngineTests.scala | 24 ------- .../kyuubi/jdbc/hive/KyuubiConnection.java | 28 ++++---- .../kyuubi/jdbc/hive/KyuubiStatement.java | 30 -------- .../kyuubi/operation/GetCurrentCatalog.scala | 32 --------- .../kyuubi/operation/GetCurrentDatabase.scala | 32 --------- .../operation/KyuubiOperationManager.scala | 17 ++--- .../kyuubi/operation/SetCurrentCatalog.scala | 33 --------- .../kyuubi/operation/SetCurrentDatabase.scala | 33 --------- ...CatalogDatabaseOperationManagerSuite.scala | 57 ++++++++++++++++ .../KyuubiOperationManagerSuite.scala | 25 ------- 24 files changed, 293 insertions(+), 298 deletions(-) create mode 100644 externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala create mode 100644 externals/kyuubi-hive-sql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala create mode 100644 externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoCatalogDatabaseOperationSuite.scala delete mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala delete mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala delete mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala delete mode 100644 kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala create mode 100644 kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala diff --git a/docs/deployment/settings.md b/docs/deployment/settings.md index c3e35ce5675..7c10b75969d 100644 --- a/docs/deployment/settings.md +++ b/docs/deployment/settings.md @@ -216,6 +216,7 @@ Key | Default | Meaning | Type | Since kyuubi.engine.hive.java.options|
<undefined>
|
The extra java options for the hive query engine
|
string
|
1.6.0
kyuubi.engine.hive.memory|
1g
|
The heap memory for the hive query engine
|
string
|
1.6.0
kyuubi.engine.initialize.sql|
SHOW DATABASES
|
SemiColon-separated list of SQL statements to be initialized in the newly created engine before queries. i.e. use `SHOW DATABASES` to eagerly active HiveClient. This configuration can not be used in JDBC url due to the limitation of Beeline/JDBC driver.
|
seq
|
1.2.0
+kyuubi.engine.operation.convert.catalog.database.enabled|
false
|
When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines
|
boolean
|
1.6.0
kyuubi.engine.operation.log.dir.root|
engine_operation_logs
|
Root directory for query operation log at engine-side.
|
string
|
1.4.0
kyuubi.engine.pool.name|
engine-pool
|
The name of engine pool.
|
string
|
1.5.0
kyuubi.engine.pool.size|
-1
|
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).
|
int
|
1.4.0
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 fcd60ac132c..6fce4c61e91 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 @@ -42,7 +42,7 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) if (catalogDatabaseOperation != null) { return catalogDatabaseOperation } diff --git a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala new file mode 100644 index 00000000000..18c52469c07 --- /dev/null +++ b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala @@ -0,0 +1,54 @@ +/* + * 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.flink.operation + +import org.apache.kyuubi.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED +import org.apache.kyuubi.engine.flink.WithFlinkSQLEngine +import org.apache.kyuubi.operation.HiveJDBCTestHelper + +class FlinkCatalogDatabaseOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { + override def withKyuubiConf: Map[String, String] = + Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") + + override protected def jdbcUrl: String = + s"jdbc:hive2://${engine.frontendServices.head.connectionUrl}/;" + + test("execute statement - set/get catalog") { + 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 - set/get database") { + 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")) + }) + } +} 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 ab65cafd1fb..41998238712 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 @@ -665,18 +665,6 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { }) } - test("execute statement - set/get catalog") { - 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 => @@ -686,18 +674,6 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { }) } - test("execute statement - set/get database") { - 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 => 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 33026ff7ef6..f99b7e94aeb 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 @@ -36,7 +36,7 @@ class HiveOperationManager() extends OperationManager("HiveOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) if (catalogDatabaseOperation != null) { return catalogDatabaseOperation } 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 00000000000..aa95c4df697 --- /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-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 e35ff2e837a..d256dc26121 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 @@ -53,7 +53,7 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n runAsync: Boolean, queryTimeout: Long): Operation = { val spark = session.asInstanceOf[SparkSessionImpl].spark - val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) if (catalogDatabaseOperation != null) { return catalogDatabaseOperation } 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 index 97380b06bd9..46208bff1e5 100644 --- 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 @@ -20,6 +20,7 @@ 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 @@ -30,7 +31,8 @@ class SparkCatalogDatabaseOperationSuite extends WithSparkSQLEngine with HiveJDB override def withKyuubiConf: Map[String, String] = Map( - "spark.sql.catalog.dummy" -> classOf[DummyCatalog].getName) + "spark.sql.catalog.dummy" -> classOf[DummyCatalog].getName, + ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") test("set/get current catalog") { withJdbcStatement() { statement => 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 811cda470e3..08d84342847 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 @@ -34,7 +34,7 @@ class TrinoOperationManager extends OperationManager("TrinoOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) + val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) if (catalogDatabaseOperation != null) { return catalogDatabaseOperation } 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 00000000000..029fbcf2a59 --- /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/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala index f5d63fdcd21..d3680c30aaa 100644 --- a/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala +++ b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperationSuite.scala @@ -738,26 +738,4 @@ class TrinoOperationSuite extends WithTrinoEngine with TrinoQueryTests { assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS) } } - - 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 f9ef694e002..5a798162f74 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 @@ -1521,4 +1521,12 @@ object KyuubiConf { .version("1.6.0") .booleanConf .createWithDefault(true) + + 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(false) } 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 190f0b044b8..39bef393d62 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,8 +25,4 @@ 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 b533167b493..8e4a3d84340 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,7 +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.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation import org.apache.kyuubi.operation.OperationState._ import org.apache.kyuubi.operation.log.LogDivertAppender @@ -159,17 +159,29 @@ abstract class OperationManager(name: String) extends AbstractService(name) { } } } + + private val PATTERN_FOR_SET_CATALOG = "_SET_CATALOG_(.+)".r + 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, - confOverlay: Map[String, String]): Operation = { - if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_CATALOG)) { - newSetCurrentCatalogOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_CATALOG)) - } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_CATALOG)) { - newGetCurrentCatalogOperation(session) - } else if (confOverlay.contains(KYUUBI_OPERATION_SET_CURRENT_DATABASE)) { - newSetCurrentDatabaseOperation(session, confOverlay(KYUUBI_OPERATION_SET_CURRENT_DATABASE)) - } else if (confOverlay.contains(KYUUBI_OPERATION_GET_CURRENT_DATABASE)) { - newGetCurrentDatabaseOperation(session) + statement: String, + conf: KyuubiConf): Operation = { + if (conf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)) { + statement match { + case PATTERN_FOR_SET_CATALOG(catalog) => + newSetCurrentCatalogOperation(session, catalog) + case _ if PATTERN_FOR_GET_CATALOG == statement => + newGetCurrentCatalogOperation(session) + case PATTERN_FOR_SET_SCHEMA(database) => + newSetCurrentDatabaseOperation(session, database) + case _ if PATTERN_FOR_GET_SCHEMA == statement.toLowerCase => + newGetCurrentDatabaseOperation(session) + case _ => + null + } } else { null } diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala index b957a8ee924..9eb4a24405b 100644 --- a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala +++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala @@ -452,28 +452,4 @@ trait HiveEngineTests extends HiveJDBCTestHelper { assert(resultSet.getString(1) === "hive.query.name=test") } } - - 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/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 b34957354ec..679bce96467 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,10 +1094,6 @@ public Statement createStatement() throws SQLException { return new KyuubiStatement(this, client, sessHandle, fetchSize); } - private KyuubiStatement createKyuubiStatement() throws SQLException { - return ((KyuubiStatement) createStatement()); - } - /* * (non-Javadoc) * @@ -1171,12 +1167,14 @@ public String getCatalog() throws SQLException { if (isClosed) { throw new SQLException("Connection is closed"); } - try (KyuubiStatement stmt = createKyuubiStatement(); - ResultSet res = stmt.executeGetCurrentCatalog()) { + try (Statement stmt = createStatement(); + ResultSet res = stmt.executeQuery("_GET_CATALOG")) { if (!res.next()) { throw new SQLException("Failed to get catalog information"); } return res.getString(1); + } catch (Exception e) { + return ""; } } @@ -1238,8 +1236,8 @@ public String getSchema() throws SQLException { if (isClosed) { throw new SQLException("Connection is closed"); } - try (KyuubiStatement stmt = createKyuubiStatement(); - ResultSet res = stmt.executeGetCurrentDatabase()) { + try (Statement stmt = createStatement(); + ResultSet res = stmt.executeQuery("SELECT current_database()")) { if (!res.next()) { throw new SQLException("Failed to get schema information"); } @@ -1527,11 +1525,17 @@ public void setAutoCommit(boolean autoCommit) throws SQLException { @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"); } - try (KyuubiStatement stmt = createKyuubiStatement()) { - stmt.executeSetCurrentCatalog(catalog); + try (Statement stmt = createStatement()) { + try { + stmt.execute("_SET_CATALOG_" + catalog); + } catch (Exception e) { + + } } } @@ -1654,8 +1658,8 @@ public void setSchema(String schema) throws SQLException { if (schema == null || schema.isEmpty()) { throw new SQLException("Schema name is null or empty"); } - try (KyuubiStatement stmt = createKyuubiStatement()) { - stmt.executeSetCurrentDatabase(schema); + try (Statement stmt = createStatement()) { + stmt.execute("use " + 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 2595c029eb5..31156f5fb28 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,36 +512,6 @@ public ResultSet executeScala(String code) throws SQLException { return resultSet; } - public void executeSetCurrentCatalog(String catalog) throws SQLException { - if (executeWithConfOverlay( - "", Collections.singletonMap("kyuubi.operation.set.current.catalog", catalog))) { - resultSet.close(); - } - } - - public ResultSet executeGetCurrentCatalog() throws SQLException { - if (!executeWithConfOverlay( - "", Collections.singletonMap("kyuubi.operation.get.current.catalog", ""))) { - throw new SQLException("The query did not generate a result set!"); - } - return resultSet; - } - - public void executeSetCurrentDatabase(String database) throws SQLException { - if (executeWithConfOverlay( - "", Collections.singletonMap("kyuubi.operation.set.current.database", database))) { - resultSet.close(); - } - } - - public ResultSet executeGetCurrentDatabase() throws SQLException { - if (!executeWithConfOverlay( - "", 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/GetCurrentCatalog.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala deleted file mode 100644 index 05045bbc984..00000000000 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentCatalog.scala +++ /dev/null @@ -1,32 +0,0 @@ -/* - * 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.operation - -import org.apache.kyuubi.config.KyuubiReservedKeys._ -import org.apache.kyuubi.session.Session - -class GetCurrentCatalog(session: Session) - extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { - - override protected def runInternal(): Unit = { - try { - val conf = Map(KYUUBI_OPERATION_GET_CURRENT_CATALOG -> "") - _remoteOpHandle = client.executeStatement("", conf, false, 0L) - } catch onError() - } -} diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala deleted file mode 100644 index 88d2bae7c32..00000000000 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCurrentDatabase.scala +++ /dev/null @@ -1,32 +0,0 @@ -/* - * 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.operation - -import org.apache.kyuubi.config.KyuubiReservedKeys._ -import org.apache.kyuubi.session.Session - -class GetCurrentDatabase(session: Session) - extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { - - override protected def runInternal(): Unit = { - try { - val conf = Map(KYUUBI_OPERATION_GET_CURRENT_DATABASE -> "") - _remoteOpHandle = client.executeStatement("", conf, false, 0L) - } catch onError() - } -} 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 78e470ef4de..7ad1c99e9cd 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 @@ -58,10 +58,6 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, confOverlay) - if (catalogDatabaseOperation != null) { - return catalogDatabaseOperation - } val operation = new ExecuteStatement(session, statement, confOverlay, runAsync, getQueryTimeout(queryTimeout)) addOperation(operation) @@ -75,24 +71,21 @@ 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 = { - val operation = new SetCurrentCatalog(session, catalog) - addOperation(operation) + null } override def newGetCurrentCatalogOperation(session: Session): Operation = { - val operation = new GetCurrentCatalog(session) - addOperation(operation) + null } override def newSetCurrentDatabaseOperation(session: Session, database: String): Operation = { - val operation = new SetCurrentDatabase(session, database) - addOperation(operation) + null } override def newGetCurrentDatabaseOperation(session: Session): Operation = { - val operation = new GetCurrentDatabase(session) - addOperation(operation) + null } override def newGetTypeInfoOperation(session: Session): Operation = { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala deleted file mode 100644 index 6a7c36eee5a..00000000000 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentCatalog.scala +++ /dev/null @@ -1,33 +0,0 @@ -/* - * 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.operation - -import org.apache.kyuubi.config.KyuubiReservedKeys._ -import org.apache.kyuubi.session.Session - -class SetCurrentCatalog(session: Session, catalog: String) - extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { - - override protected def runInternal(): Unit = { - try { - val conf = Map(KYUUBI_OPERATION_SET_CURRENT_CATALOG -> catalog) - _remoteOpHandle = client.executeStatement("", conf, false, 0L) - setHasResultSet(_remoteOpHandle.isHasResultSet) - } catch onError() - } -} diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala deleted file mode 100644 index 0da3f1b561b..00000000000 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/SetCurrentDatabase.scala +++ /dev/null @@ -1,33 +0,0 @@ -/* - * 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.operation - -import org.apache.kyuubi.config.KyuubiReservedKeys._ -import org.apache.kyuubi.session.Session - -class SetCurrentDatabase(session: Session, database: String) - extends KyuubiOperation(OperationType.EXECUTE_STATEMENT, session) { - - override protected def runInternal(): Unit = { - try { - val conf = Map(KYUUBI_OPERATION_SET_CURRENT_DATABASE -> database) - _remoteOpHandle = client.executeStatement("", conf, false, 0L) - setHasResultSet(_remoteOpHandle.isHasResultSet) - } catch onError() - } -} diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala new file mode 100644 index 00000000000..1b86828c24e --- /dev/null +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala @@ -0,0 +1,57 @@ +/* + * 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.operation + +import java.sql.SQLException + +import org.apache.kyuubi.WithKyuubiServer +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.config.KyuubiConf._ + +class KyuubiCatalogDatabaseOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelper { + override protected val conf: KyuubiConf = { + KyuubiConf().set(OPERATION_QUERY_TIMEOUT.key, "PT1S") + .set(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key, "true") + } + + test("set/get current catalog") { + withJdbcStatement() { statement => + val catalog = statement.getConnection.getCatalog + assert(catalog == "spark_catalog") + // The server starts the spark engine without other catalogs + val e = intercept[SQLException] { + statement.getConnection.setCatalog("dummy_catalog") + statement.getConnection.getCatalog + } + assert(e.getMessage.contains("dummy_catalog")) + } + } + + 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") + } + } + + override protected def jdbcUrl: String = getJdbcUrl +} 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 abae65d0d55..e7082463b98 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 @@ -17,7 +17,6 @@ package org.apache.kyuubi.operation -import java.sql.SQLException import java.sql.SQLTimeoutException import org.apache.kyuubi.WithKyuubiServer @@ -61,29 +60,5 @@ class KyuubiOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelp } } - test("set/get current catalog") { - withJdbcStatement() { statement => - val catalog = statement.getConnection.getCatalog - assert(catalog == "spark_catalog") - // The server starts the spark engine without other catalogs - val e = intercept[SQLException] { - statement.getConnection.setCatalog("dummy_catalog") - statement.getConnection.getCatalog - } - assert(e.getMessage.contains("dummy_catalog")) - } - } - - 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") - } - } - override protected def jdbcUrl: String = getJdbcUrl } From ff9078abc56a1b30a60fed52a76ff89ee8309336 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 02:27:48 +0800 Subject: [PATCH 08/13] fix UT --- .../operation/FlinkCatalogDatabaseOperationSuite.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala index 18c52469c07..887a13998e1 100644 --- a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala +++ b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala @@ -17,13 +17,15 @@ package org.apache.kyuubi.engine.flink.operation -import org.apache.kyuubi.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED +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.operation.HiveJDBCTestHelper class FlinkCatalogDatabaseOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { override def withKyuubiConf: Map[String, String] = - Map(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") + Map(OPERATION_PLAN_ONLY_MODE.key -> NONE.toString, + ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") override protected def jdbcUrl: String = s"jdbc:hive2://${engine.frontendServices.head.connectionUrl}/;" From e55bb2e2143fc642dfae899e1452eb44f18e58f4 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 02:57:23 +0800 Subject: [PATCH 09/13] style --- .../flink/operation/FlinkCatalogDatabaseOperationSuite.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala index 887a13998e1..e831394f53c 100644 --- a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala +++ b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala @@ -24,7 +24,8 @@ import org.apache.kyuubi.operation.HiveJDBCTestHelper class FlinkCatalogDatabaseOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { override def withKyuubiConf: Map[String, String] = - Map(OPERATION_PLAN_ONLY_MODE.key -> NONE.toString, + Map( + OPERATION_PLAN_ONLY_MODE.key -> NONE.toString, ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") override protected def jdbcUrl: String = From 3495b2800b64c6bd3d62d851ec952f21512ac0b8 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 03:04:38 +0800 Subject: [PATCH 10/13] fix jdbc --- .../jdbc/operation/JdbcOperationManager.scala | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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 d4455e93cfe..c9700870b2d 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() + } } From 8bc40b9fc46c846f290784f9ecbdcab41c576d54 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 04:15:45 +0800 Subject: [PATCH 11/13] fix UT --- .../KyuubiCatalogDatabaseOperationManagerSuite.scala | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala index 1b86828c24e..8729e0d39d0 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala @@ -34,11 +34,9 @@ class KyuubiCatalogDatabaseOperationManagerSuite extends WithKyuubiServer with H val catalog = statement.getConnection.getCatalog assert(catalog == "spark_catalog") // The server starts the spark engine without other catalogs - val e = intercept[SQLException] { - statement.getConnection.setCatalog("dummy_catalog") - statement.getConnection.getCatalog - } - assert(e.getMessage.contains("dummy_catalog")) + statement.getConnection.setCatalog("dummy_catalog") + val changedCatalog = statement.getConnection.getCatalog + assert(changedCatalog == "spark_catalog") } } From 2acf0a453aa4d8addda049b95184f8da441199e3 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 10:46:43 +0800 Subject: [PATCH 12/13] unused import --- .../operation/KyuubiCatalogDatabaseOperationManagerSuite.scala | 2 -- 1 file changed, 2 deletions(-) diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala index 8729e0d39d0..4e99be1f0f4 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala @@ -17,8 +17,6 @@ package org.apache.kyuubi.operation -import java.sql.SQLException - import org.apache.kyuubi.WithKyuubiServer import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf._ From 670516443944f356dd3a0694e2daa4071fdf6599 Mon Sep 17 00:00:00 2001 From: sychen Date: Thu, 26 May 2022 15:47:54 +0800 Subject: [PATCH 13/13] use placeholders and confOverlay --- docs/deployment/settings.md | 2 +- .../operation/FlinkSQLOperationManager.scala | 15 +++-- .../FlinkCatalogDatabaseOperationSuite.scala | 57 ------------------- .../flink/operation/FlinkOperationSuite.scala | 38 ++++++++++++- .../hive/operation/HiveOperationManager.scala | 9 ++- .../operation/SparkSQLOperationManager.scala | 11 +++- .../engine/spark/shim/CatalogShim_v3_0.scala | 7 ++- .../operation/TrinoOperationManager.scala | 10 ++-- .../org/apache/kyuubi/config/KyuubiConf.scala | 2 +- .../kyuubi/config/KyuubiReservedKeys.scala | 4 ++ .../kyuubi/operation/OperationManager.scala | 31 +++++----- .../kyuubi/jdbc/hive/KyuubiConnection.java | 26 ++++----- .../kyuubi/jdbc/hive/KyuubiStatement.java | 30 ++++++++++ ...CatalogDatabaseOperationManagerSuite.scala | 53 ----------------- .../KyuubiOperationManagerSuite.scala | 32 ++++++++++- 15 files changed, 167 insertions(+), 160 deletions(-) delete mode 100644 externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala delete mode 100644 kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala diff --git a/docs/deployment/settings.md b/docs/deployment/settings.md index 2756692593a..d1d88d97229 100644 --- a/docs/deployment/settings.md +++ b/docs/deployment/settings.md @@ -223,7 +223,7 @@ Key | Default | Meaning | Type | Since kyuubi.engine.jdbc.connection.user|
<undefined>
|
The user is used for connecting to server
|
string
|
1.6.0
kyuubi.engine.jdbc.driver.class|
<undefined>
|
The driver class for jdbc engine connection
|
string
|
1.6.0
kyuubi.engine.jdbc.type|
<undefined>
|
The short name of jdbc type
|
string
|
1.6.0
-kyuubi.engine.operation.convert.catalog.database.enabled|
false
|
When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines
|
boolean
|
1.6.0
+kyuubi.engine.operation.convert.catalog.database.enabled|
true
|
When set to true, The engine converts the JDBC methods of set/get Catalog and set/get Schema to the implementation of different engines
|
boolean
|
1.6.0
kyuubi.engine.operation.log.dir.root|
engine_operation_logs
|
Root directory for query operation log at engine-side.
|
string
|
1.4.0
kyuubi.engine.pool.name|
engine-pool
|
The name of engine pool.
|
string
|
1.5.0
kyuubi.engine.pool.size|
-1
|
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).
|
int
|
1.4.0
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 6fce4c61e91..056e7844f97 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,17 +36,24 @@ 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, confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) - if (catalogDatabaseOperation != null) { - return catalogDatabaseOperation - } 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) diff --git a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala b/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala deleted file mode 100644 index e831394f53c..00000000000 --- a/externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkCatalogDatabaseOperationSuite.scala +++ /dev/null @@ -1,57 +0,0 @@ -/* - * 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.flink.operation - -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.operation.HiveJDBCTestHelper - -class FlinkCatalogDatabaseOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper { - override def withKyuubiConf: Map[String, String] = - Map( - OPERATION_PLAN_ONLY_MODE.key -> NONE.toString, - ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key -> "true") - - override protected def jdbcUrl: String = - s"jdbc:hive2://${engine.frontendServices.head.connectionUrl}/;" - - test("execute statement - set/get catalog") { - 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 - set/get database") { - 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")) - }) - } -} 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 f96cc61bcfd..90dd0025725 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/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 f99b7e94aeb..8cd7a3112ce 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,9 +37,11 @@ class HiveOperationManager() extends OperationManager("HiveOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) - if (catalogDatabaseOperation != null) { - return catalogDatabaseOperation + 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) 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 d256dc26121..30ad1c203b8 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,9 +55,12 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n runAsync: Boolean, queryTimeout: Long): Operation = { val spark = session.asInstanceOf[SparkSessionImpl].spark - val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) - if (catalogDatabaseOperation != null) { - return catalogDatabaseOperation + 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, 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 4dd3a230b95..e42d2cf330e 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 @@ -55,7 +55,12 @@ class CatalogShim_v3_0 extends CatalogShim_v2_4 { } override def setCurrentCatalog(spark: SparkSession, catalog: String): Unit = { - spark.sessionState.catalogManager.setCurrentCatalog(catalog) + // 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 = { 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 08d84342847..63140bd7fc4 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,9 +34,11 @@ class TrinoOperationManager extends OperationManager("TrinoOperationManager") { confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): Operation = { - val catalogDatabaseOperation = processCatalogDatabase(session, statement, getConf) - if (catalogDatabaseOperation != null) { - return catalogDatabaseOperation + 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 = 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 b62b1551b7c..863734ed324 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 @@ -1578,5 +1578,5 @@ object KyuubiConf { "and set/get Schema to the implementation of different engines") .version("1.6.0") .booleanConf - .createWithDefault(false) + .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 39bef393d62..190f0b044b8 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 8e4a3d84340..d5b45623461 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,7 +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.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED +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 @@ -160,7 +160,7 @@ abstract class OperationManager(name: String) extends AbstractService(name) { } } - private val PATTERN_FOR_SET_CATALOG = "_SET_CATALOG_(.+)".r + 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()" @@ -168,20 +168,19 @@ abstract class OperationManager(name: String) extends AbstractService(name) { final def processCatalogDatabase( session: Session, statement: String, - conf: KyuubiConf): Operation = { - if (conf.get(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED)) { - statement match { - case PATTERN_FOR_SET_CATALOG(catalog) => - newSetCurrentCatalogOperation(session, catalog) - case _ if PATTERN_FOR_GET_CATALOG == statement => - newGetCurrentCatalogOperation(session) - case PATTERN_FOR_SET_SCHEMA(database) => - newSetCurrentDatabaseOperation(session, database) - case _ if PATTERN_FOR_GET_SCHEMA == statement.toLowerCase => - newGetCurrentDatabaseOperation(session) - case _ => - null - } + 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-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 679bce96467..5aa94e9471b 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 Statement createStatement() throws SQLException { return new KyuubiStatement(this, client, sessHandle, fetchSize); } + private KyuubiStatement createKyuubiStatement() throws SQLException { + return ((KyuubiStatement) createStatement()); + } + /* * (non-Javadoc) * @@ -1167,8 +1171,8 @@ public String getCatalog() throws SQLException { if (isClosed) { throw new SQLException("Connection is closed"); } - try (Statement stmt = createStatement(); - ResultSet res = stmt.executeQuery("_GET_CATALOG")) { + try (KyuubiStatement stmt = createKyuubiStatement(); + ResultSet res = stmt.executeGetCurrentCatalog("_GET_CATALOG")) { if (!res.next()) { throw new SQLException("Failed to get catalog information"); } @@ -1236,8 +1240,8 @@ public String getSchema() throws SQLException { 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"); } @@ -1525,17 +1529,13 @@ public void setAutoCommit(boolean autoCommit) throws SQLException { @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"); } - try (Statement stmt = createStatement()) { - try { - stmt.execute("_SET_CATALOG_" + catalog); - } catch (Exception e) { + try (KyuubiStatement stmt = createKyuubiStatement()) { + stmt.executeSetCurrentCatalog("_SET_CATALOG", catalog); + } catch (SQLException e) { - } } } @@ -1658,8 +1658,8 @@ public void setSchema(String schema) throws SQLException { 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 31156f5fb28..4f13a187ad5 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 ResultSet executeScala(String code) throws SQLException { 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/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala deleted file mode 100644 index 4e99be1f0f4..00000000000 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiCatalogDatabaseOperationManagerSuite.scala +++ /dev/null @@ -1,53 +0,0 @@ -/* - * 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.operation - -import org.apache.kyuubi.WithKyuubiServer -import org.apache.kyuubi.config.KyuubiConf -import org.apache.kyuubi.config.KyuubiConf._ - -class KyuubiCatalogDatabaseOperationManagerSuite extends WithKyuubiServer with HiveJDBCTestHelper { - override protected val conf: KyuubiConf = { - KyuubiConf().set(OPERATION_QUERY_TIMEOUT.key, "PT1S") - .set(ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key, "true") - } - - test("set/get current catalog") { - 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") { - 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 -} 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 e7082463b98..1bd750e6aed 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 }