diff --git a/docs/deployment/settings.md b/docs/deployment/settings.md index 4a24499be..4fa8c73d9 100644 --- a/docs/deployment/settings.md +++ b/docs/deployment/settings.md @@ -154,8 +154,10 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co | 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.extra.classpath | <undefined> | The extra classpath for the JDBC query engine, for configuring the location of the JDBC driver and etc. | string | 1.6.0 | +| kyuubi.engine.jdbc.initialize.sql | SELECT 1 | SemiColon-separated list of SQL statements to be initialized in the newly created engine before queries. i.e. use `SELECT 1` to eagerly active JDBCClient. | seq | 1.8.0 | | kyuubi.engine.jdbc.java.options | <undefined> | The extra Java options for the JDBC query engine | string | 1.6.0 | | kyuubi.engine.jdbc.memory | 1g | The heap memory for the JDBC query engine | string | 1.6.0 | +| kyuubi.engine.jdbc.session.initialize.sql || SemiColon-separated list of SQL statements to be initialized in the newly created engine session before queries. | seq | 1.8.0 | | kyuubi.engine.jdbc.type | <undefined> | The short name of JDBC type | string | 1.6.0 | | kyuubi.engine.kubernetes.submit.timeout | PT30S | The engine submit timeout for Kubernetes application. | duration | 1.7.2 | | 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 | diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala index 618098f31..6e0647f6c 100644 --- a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala +++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala @@ -19,7 +19,9 @@ package org.apache.kyuubi.engine.jdbc import org.apache.kyuubi.{Logging, Utils} import org.apache.kyuubi.Utils.{addShutdownHook, JDBC_ENGINE_SHUTDOWN_PRIORITY} import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.config.KyuubiConf.ENGINE_JDBC_INITIALIZE_SQL import org.apache.kyuubi.engine.jdbc.JdbcSQLEngine.currentEngine +import org.apache.kyuubi.engine.jdbc.util.KyuubiJdbcUtils import org.apache.kyuubi.ha.HighAvailabilityConf.HA_ZK_CONN_RETRY_POLICY import org.apache.kyuubi.ha.client.RetryPolicies import org.apache.kyuubi.service.Serverable @@ -71,6 +73,8 @@ object JdbcSQLEngine extends Logging { kyuubiConf.setIfMissing(HA_ZK_CONN_RETRY_POLICY, RetryPolicies.N_TIME.toString) startEngine() + + KyuubiJdbcUtils.initializeJdbcSession(kyuubiConf, kyuubiConf.get(ENGINE_JDBC_INITIALIZE_SQL)) } catch { case t: Throwable if currentEngine.isDefined => currentEngine.foreach { engine => diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionImpl.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionImpl.scala index f8cd40412..5acae8f24 100644 --- a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionImpl.scala +++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionImpl.scala @@ -23,8 +23,10 @@ import scala.util.{Failure, Success, Try} import org.apache.hive.service.rpc.thrift.{TGetInfoType, TGetInfoValue, TProtocolVersion} import org.apache.kyuubi.KyuubiSQLException +import org.apache.kyuubi.config.KyuubiConf.ENGINE_JDBC_SESSION_INITIALIZE_SQL import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider +import org.apache.kyuubi.engine.jdbc.util.KyuubiJdbcUtils import org.apache.kyuubi.session.{AbstractSession, SessionHandle, SessionManager} class JdbcSessionImpl( @@ -51,6 +53,10 @@ class JdbcSessionImpl( sessionConnection = ConnectionProvider.create(kyuubiConf) databaseMetaData = sessionConnection.getMetaData } + KyuubiJdbcUtils.initializeJdbcSession( + kyuubiConf, + sessionConnection, + kyuubiConf.get(ENGINE_JDBC_SESSION_INITIALIZE_SQL)) super.open() info(s"The jdbc session is started.") } diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/util/KyuubiJdbcUtils.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/util/KyuubiJdbcUtils.scala new file mode 100644 index 000000000..7107045ff --- /dev/null +++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/util/KyuubiJdbcUtils.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.engine.jdbc.util + +import java.sql.Connection + +import org.apache.kyuubi.{KyuubiSQLException, Logging} +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider +import org.apache.kyuubi.engine.jdbc.dialect.{JdbcDialect, JdbcDialects} +import org.apache.kyuubi.util.JdbcUtils + +object KyuubiJdbcUtils extends Logging { + + def initializeJdbcSession(kyuubiConf: KyuubiConf, initializationSQLs: Seq[String]): Unit = { + JdbcUtils.withCloseable(ConnectionProvider.create(kyuubiConf)) { connection => + initializeJdbcSession(kyuubiConf, connection, initializationSQLs) + } + } + + def initializeJdbcSession( + kyuubiConf: KyuubiConf, + connection: Connection, + initializationSQLs: Seq[String]): Unit = { + if (initializationSQLs == null || initializationSQLs.isEmpty) { + return + } + try { + val dialect: JdbcDialect = JdbcDialects.get(kyuubiConf) + JdbcUtils.withCloseable(dialect.createStatement(connection)) { statement => + initializationSQLs.foreach { sql => + debug(s"Execute initialization sql: $sql") + statement.execute(sql) + } + } + } catch { + case e: Exception => + error("Failed to execute initialization sql.", e) + throw KyuubiSQLException(e) + } + } +} 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 36fd27285..d51a5f5c3 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 @@ -2585,6 +2585,24 @@ object KyuubiConf { .stringConf .createOptional + val ENGINE_JDBC_INITIALIZE_SQL: ConfigEntry[Seq[String]] = + buildConf("kyuubi.engine.jdbc.initialize.sql") + .doc("SemiColon-separated list of SQL statements to be initialized in the newly created " + + "engine before queries. i.e. use `SELECT 1` to eagerly active JDBCClient.") + .version("1.8.0") + .stringConf + .toSequence(";") + .createWithDefaultString("SELECT 1") + + val ENGINE_JDBC_SESSION_INITIALIZE_SQL: ConfigEntry[Seq[String]] = + buildConf("kyuubi.engine.jdbc.session.initialize.sql") + .doc("SemiColon-separated list of SQL statements to be initialized in the newly created " + + "engine session before queries.") + .version("1.8.0") + .stringConf + .toSequence(";") + .createWithDefault(Nil) + 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 " +