[KYUUBI #6253] Support running JDBC engine on YARN AM

# 🔍 Description ## Issue References 🔗 This pull request fixes #6253 ## Describe Your Solution 🔧 Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change. ## Types of changes 🔖 - [ ] Bugfix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) ## Test Plan 🧪 #### Behavior Without This Pull Request ⚰️ #### Behavior With This Pull Request 🎉 <img width="1251" alt="image" src="https://github.com/apache/kyuubi/assets/18713676/b654a300-8c79-4461-9fba-4ad1c913accc"> #### Related Unit Tests --- # Checklist 📝 - [ ] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html) **Be nice. Be informative.** Closes #6275 from lsm1/branch-jdbc-engine-on-yarn. Closes #6253 5ed4af0 [senmiaoliu] fix style 86e0326 [senmiaoliu] fix style b3e1144 [senmiaoliu] fix style bf33de3 [senmiaoliu] fix style c384049 [senmiaoliu] jdbc engine on yarn Authored-by: senmiaoliu <[email protected]> Signed-off-by: Shaoyun Chen <[email protected]>
apache · Apr 15, 2024 · 9aa5269 · 9aa5269
1 parent 3133b59
commit 9aa5269
Show file tree

Hide file tree

Showing 24 changed files with 558 additions and 72 deletions.
diff --git a/docs/configuration/settings.md b/docs/configuration/settings.md
@@ -161,6 +161,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
 | kyuubi.engine.jdbc.connection.provider                   | &lt;undefined&gt;         | A JDBC connection provider plugin for the Kyuubi Server to establish a connection to the JDBC URL. The configuration value should be a subclass of `org.apache.kyuubi.engine.jdbc.connection.JdbcConnectionProvider`. Kyuubi provides the following built-in implementations: <li>doris: For establishing Doris connections.</li> <li>mysql: For establishing MySQL connections.</li> <li>phoenix: For establishing Phoenix connections.</li> <li>postgresql: For establishing PostgreSQL connections.</li><li>starrocks: For establishing StarRocks connections.</li><li>impala: For establishing Impala connections.</li><li>clickhouse: For establishing clickhouse connections.</li>                                                                                                                                                                                                                                                                                                                                                                        | string   | 1.6.0  |
 | kyuubi.engine.jdbc.connection.url                        | &lt;undefined&gt;         | The server url that engine will connect to                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      | string   | 1.6.0  |
 | kyuubi.engine.jdbc.connection.user                       | &lt;undefined&gt;         | The user is used for connecting to server                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       | string   | 1.6.0  |
+| kyuubi.engine.jdbc.deploy.mode                           | LOCAL                     | Configures the jdbc engine deploy mode, The value can be 'local', 'yarn'. In local mode, the engine operates on the same node as the KyuubiServer. In YARN mode, the engine runs within the Application Master (AM) container of YARN.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          | string   | 1.10.0 |
 | kyuubi.engine.jdbc.driver.class                          | &lt;undefined&gt;         | The driver class for JDBC engine connection                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     | string   | 1.6.0  |
 | kyuubi.engine.jdbc.extra.classpath                       | &lt;undefined&gt;         | 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.fetch.size                            | 1000                      | The fetch size of JDBC engine                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | int      | 1.9.0  |

diff --git a/...ql-engine/src/main/scala/org/apache/kyuubi/engine/hive/deploy/HiveYarnModeSubmitter.scala b/...ql-engine/src/main/scala/org/apache/kyuubi/engine/hive/deploy/HiveYarnModeSubmitter.scala
@@ -26,6 +26,7 @@ import org.apache.kyuubi.Utils
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.ENGINE_HIVE_EXTRA_CLASSPATH
 import org.apache.kyuubi.engine.deploy.yarn.EngineYarnModeSubmitter
+import org.apache.kyuubi.engine.deploy.yarn.EngineYarnModeSubmitter.KYUUBI_ENGINE_DEPLOY_YARN_MODE_HIVE_CONF_KEY
 import org.apache.kyuubi.engine.hive.HiveSQLEngine
 
 object HiveYarnModeSubmitter extends EngineYarnModeSubmitter {
@@ -60,18 +61,10 @@ object HiveYarnModeSubmitter extends EngineYarnModeSubmitter {
     jars.toSeq
   }
 
-  private[hive] def parseClasspath(classpath: String, jars: ListBuffer[File]): Unit = {
-    classpath.split(":").filter(_.nonEmpty).foreach { cp =>
-      if (cp.endsWith("/*")) {
-        val dir = cp.substring(0, cp.length - 2)
-        new File(dir) match {
-          case f if f.isDirectory =>
-            f.listFiles().filter(_.getName.endsWith(".jar")).foreach(jars += _)
-          case _ =>
-        }
-      } else {
-        jars += new File(cp)
-      }
-    }
+  override def listConfFiles(): Seq[File] = {
+    // respect the following priority loading configuration, and distinct files
+    // hive configuration -> hadoop configuration -> yarn configuration
+    val hiveConf = kyuubiConf.getOption(KYUUBI_ENGINE_DEPLOY_YARN_MODE_HIVE_CONF_KEY)
+    listDistinctFiles(hiveConf.get) ++ super.listConfFiles()
   }
 }
diff --git a/...gine/src/test/scala/org/apache/kyuubi/engine/hive/deploy/HiveYarnModeSubmitterSuite.scala b/...gine/src/test/scala/org/apache/kyuubi/engine/hive/deploy/HiveYarnModeSubmitterSuite.scala
diff --git a/...rnals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala b/...rnals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcSQLEngine.scala
@@ -16,10 +16,15 @@
  */
 package org.apache.kyuubi.engine.jdbc
 
+import java.security.PrivilegedExceptionAction
+
+import org.apache.hadoop.security.UserGroupInformation
+
 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.config.{KyuubiConf, KyuubiReservedKeys}
+import org.apache.kyuubi.config.KyuubiConf.{ENGINE_JDBC_DEPLOY_MODE, ENGINE_JDBC_INITIALIZE_SQL, ENGINE_KEYTAB, ENGINE_PRINCIPAL}
+import org.apache.kyuubi.engine.deploy.DeployMode
 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
@@ -38,6 +43,7 @@ class JdbcSQLEngine extends Serverable("JdbcSQLEngine") {
     // Start engine self-terminating checker after all services are ready and it can be reached by
     // all servers in engine spaces.
     backendService.sessionManager.startTerminatingChecker(() => {
+      selfExited = true
       currentEngine.foreach(_.stop())
     })
   }
@@ -71,10 +77,32 @@ object JdbcSQLEngine extends Logging {
       Utils.fromCommandLineArgs(args, kyuubiConf)
       kyuubiConf.setIfMissing(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)
       kyuubiConf.setIfMissing(HA_ZK_CONN_RETRY_POLICY, RetryPolicies.N_TIME.toString)
+      val proxyUser = kyuubiConf.getOption(KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY)
+      require(proxyUser.isDefined, s"${KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY} is not set")
+      val realUser = UserGroupInformation.getLoginUser
+      val principal = kyuubiConf.get(ENGINE_PRINCIPAL)
+      val keytab = kyuubiConf.get(ENGINE_KEYTAB)
 
-      startEngine()
+      val ugi = DeployMode.withName(kyuubiConf.get(ENGINE_JDBC_DEPLOY_MODE)) match {
+        case DeployMode.LOCAL
+            if UserGroupInformation.isSecurityEnabled && principal.isDefined && keytab.isDefined =>
+          UserGroupInformation.loginUserFromKeytab(principal.get, keytab.get)
+          UserGroupInformation.getCurrentUser
+        case DeployMode.LOCAL if proxyUser.get != realUser.getShortUserName =>
+          kyuubiConf.unset(KyuubiReservedKeys.KYUUBI_ENGINE_CREDENTIALS_KEY)
+          UserGroupInformation.createProxyUser(proxyUser.get, realUser)
+        case _ =>
+          UserGroupInformation.getCurrentUser
+      }
 
-      KyuubiJdbcUtils.initializeJdbcSession(kyuubiConf, kyuubiConf.get(ENGINE_JDBC_INITIALIZE_SQL))
+      ugi.doAs(new PrivilegedExceptionAction[Unit] {
+        override def run(): Unit = {
+          startEngine()
+          KyuubiJdbcUtils.initializeJdbcSession(
+            kyuubiConf,
+            kyuubiConf.get(ENGINE_JDBC_INITIALIZE_SQL))
+        }
+      })
     } catch {
       case t: Throwable if currentEngine.isDefined =>
         currentEngine.foreach { engine =>

diff --git a/...bc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/deploy/JdbcYarnModeSubmitter.scala b/...bc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/deploy/JdbcYarnModeSubmitter.scala
@@ -0,0 +1,60 @@
+/*
+ * 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.deploy
+
+import java.io.File
+
+import scala.collection.mutable.ListBuffer
+
+import org.apache.hadoop.security.UserGroupInformation
+
+import org.apache.kyuubi.Utils
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.ENGINE_JDBC_EXTRA_CLASSPATH
+import org.apache.kyuubi.engine.deploy.yarn.EngineYarnModeSubmitter
+import org.apache.kyuubi.engine.jdbc.JdbcSQLEngine
+
+object JdbcYarnModeSubmitter extends EngineYarnModeSubmitter {
+
+  def main(args: Array[String]): Unit = {
+    Utils.fromCommandLineArgs(args, kyuubiConf)
+
+    if (UserGroupInformation.isSecurityEnabled) {
+      require(
+        kyuubiConf.get(KyuubiConf.ENGINE_PRINCIPAL).isDefined
+          && kyuubiConf.get(KyuubiConf.ENGINE_KEYTAB).isDefined,
+        s"${KyuubiConf.ENGINE_PRINCIPAL.key} and " +
+          s"${KyuubiConf.ENGINE_KEYTAB.key} must be set when submit " +
+          s"${JdbcSQLEngine.getClass.getSimpleName.stripSuffix("$")} to YARN")
+    }
+    run()
+  }
+
+  override var engineType: String = "jdbc"
+
+  override def engineMainClass(): String = JdbcSQLEngine.getClass.getName
+
+  /**
+   * Jar list for the JDBC engine.
+   */
+  override def engineExtraJars(): Seq[File] = {
+    val extraCp = kyuubiConf.get(ENGINE_JDBC_EXTRA_CLASSPATH)
+    val jars = new ListBuffer[File]
+    extraCp.foreach(cp => parseClasspath(cp, jars))
+    jars.toSeq
+  }
+}
diff --git a/...jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionManager.scala b/...jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/session/JdbcSessionManager.scala
@@ -61,6 +61,9 @@ class JdbcSessionManager(name: String)
   }
 
   private def stopSession(): Unit = {
-    JdbcSQLEngine.currentEngine.foreach(_.stop())
+    JdbcSQLEngine.currentEngine.foreach { engine =>
+      engine.selfExited = true
+      engine.stop()
+    }
   }
 }
diff --git a/...engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/WithClickHouseEngine.scala b/...engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/WithClickHouseEngine.scala
@@ -17,6 +17,7 @@
 package org.apache.kyuubi.engine.jdbc.clickhouse
 
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithClickHouseEngine extends WithJdbcEngine with WithClickHouseContainer {
@@ -29,6 +30,7 @@ trait WithClickHouseEngine extends WithJdbcEngine with WithClickHouseContainer {
       ENGINE_JDBC_CONNECTION_PASSWORD.key -> container.password,
       ENGINE_TYPE.key -> "jdbc",
       ENGINE_JDBC_SHORT_NAME.key -> "clickhouse",
+      KYUUBI_SESSION_USER_KEY -> "kyuubi",
       ENGINE_JDBC_DRIVER_CLASS.key -> container.driverClassName)
   }
 }
diff --git a/...uubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/doris/WithDorisEngine.scala b/...uubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/doris/WithDorisEngine.scala
@@ -17,6 +17,7 @@
 package org.apache.kyuubi.engine.jdbc.doris
 
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithDorisEngine extends WithJdbcEngine with WithDorisContainer {
@@ -28,5 +29,6 @@ trait WithDorisEngine extends WithJdbcEngine with WithDorisContainer {
     ENGINE_JDBC_CONNECTION_PASSWORD.key -> "",
     ENGINE_TYPE.key -> "jdbc",
     ENGINE_JDBC_SHORT_NAME.key -> "doris",
+    KYUUBI_SESSION_USER_KEY -> "kyuubi",
     ENGINE_JDBC_DRIVER_CLASS.key -> "com.mysql.cj.jdbc.Driver")
 }
diff --git a/...bi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/impala/WithImpalaEngine.scala b/...bi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/impala/WithImpalaEngine.scala
@@ -17,6 +17,7 @@
 package org.apache.kyuubi.engine.jdbc.impala
 
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithImpalaEngine extends WithJdbcEngine with WithImpalaContainer {
@@ -26,5 +27,6 @@ trait WithImpalaEngine extends WithJdbcEngine with WithImpalaContainer {
     ENGINE_JDBC_CONNECTION_URL.key -> hiveServerJdbcUrl,
     ENGINE_TYPE.key -> "jdbc",
     ENGINE_JDBC_SHORT_NAME.key -> "impala",
+    KYUUBI_SESSION_USER_KEY -> "kyuubi",
     ENGINE_JDBC_DRIVER_CLASS.key -> ImpalaConnectionProvider.driverClass)
 }
diff --git a/...uubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/mysql/WithMySQLEngine.scala b/...uubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/mysql/WithMySQLEngine.scala
@@ -21,6 +21,7 @@ import com.dimafeng.testcontainers.scalatest.TestContainerForAll
 import org.testcontainers.utility.DockerImageName
 
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithMySQLEngine extends WithJdbcEngine with TestContainerForAll {
@@ -40,6 +41,7 @@ trait WithMySQLEngine extends WithJdbcEngine with TestContainerForAll {
       ENGINE_JDBC_CONNECTION_PASSWORD.key -> mysqlContainer.password,
       ENGINE_TYPE.key -> "jdbc",
       ENGINE_JDBC_SHORT_NAME.key -> "mysql",
+      KYUUBI_SESSION_USER_KEY -> "kyuubi",
       ENGINE_JDBC_DRIVER_CLASS.key -> "com.mysql.cj.jdbc.Driver")
   }
 }
diff --git a/...-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/phoenix/WithPhoenixEngine.scala b/...-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/phoenix/WithPhoenixEngine.scala
@@ -17,6 +17,7 @@
 package org.apache.kyuubi.engine.jdbc.phoenix
 
 import org.apache.kyuubi.config.KyuubiConf._
+import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithPhoenixEngine extends WithJdbcEngine with WithPhoenixContainer {
@@ -32,6 +33,7 @@ trait WithPhoenixEngine extends WithJdbcEngine with WithPhoenixContainer {
     ENGINE_JDBC_CONNECTION_PASSWORD.key -> "",
     ENGINE_TYPE.key -> "jdbc",
     ENGINE_JDBC_SHORT_NAME.key -> "phoenix",
+    KYUUBI_SESSION_USER_KEY -> "kyuubi",
     ENGINE_JDBC_DRIVER_CLASS.key -> "org.apache.phoenix.queryserver.client.Driver")
 
   private def getConnectString: String = s"$jdbcUrlPrefix=http://$queryServerUrl;$serialization"