From 33bdbbdab068720082c3ea0beedbc2d1d1031801 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Fri, 8 Jan 2021 23:43:14 +0800 Subject: [PATCH] Support launch SparkProcessBuilder with keytab and principal (#269) * Support launch SparkProcessBuilder with keytab and principal * fix tests --- .../spark/SparkSQLEngineListenerSuite.scala | 2 +- kyuubi-main/pom.xml | 10 ++++++ .../engine/spark/SparkProcessBuilder.scala | 33 +++++++++++++++++-- .../spark/SparkProcessBuilderSuite.scala | 32 ++++++++++++++++-- 4 files changed, 71 insertions(+), 6 deletions(-) diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala index 0c2a956a7e6..d65734e76ac 100644 --- a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala +++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/SparkSQLEngineListenerSuite.scala @@ -33,7 +33,7 @@ class SparkSQLEngineListenerSuite extends KyuubiFunSuite { test("application end") { val spark = SparkSession - .builder().master("local").getOrCreate() + .builder().master("local").config("spark.ui.port", "0").getOrCreate() val engine = new SparkSQLEngine(spark) engine.initialize(KyuubiConf()) diff --git a/kyuubi-main/pom.xml b/kyuubi-main/pom.xml index c3d50423d9c..9ebb2eca493 100644 --- a/kyuubi-main/pom.xml +++ b/kyuubi-main/pom.xml @@ -65,6 +65,16 @@ hive-jdbc test + + + org.apache.hadoop + hadoop-minikdc + + + + org.apache.directory.server + apacheds-service + diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala index 213460b85c7..3435a2c9d24 100644 --- a/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala +++ b/kyuubi-main/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala @@ -17,10 +17,13 @@ package org.apache.kyuubi.engine.spark +import java.io.IOException import java.nio.file.{Files, Path, Paths} import scala.collection.mutable.ArrayBuffer +import org.apache.hadoop.security.UserGroupInformation + import org.apache.kyuubi._ import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf.ENGINE_SPARK_MAIN_RESOURCE @@ -108,8 +111,11 @@ class SparkProcessBuilder( buffer += CONF buffer += s"$k=$v" } - buffer += PROXY_USER - buffer += proxyUser + // iff the keytab is specified, PROXY_USER is not supported + if (!useKeytab()) { + buffer += PROXY_USER + buffer += proxyUser + } mainResource.foreach { r => buffer += r } @@ -119,11 +125,32 @@ class SparkProcessBuilder( override def toString: String = commands.mkString(" ") override protected def module: String = "kyuubi-spark-sql-engine" + + private def useKeytab(): Boolean = { + val principal = conf.get(PRINCIPAL) + val keytab = conf.get(KEYTAB) + if (principal.isEmpty || keytab.isEmpty) { + false + } else { + try { + val ugi = UserGroupInformation + .loginUserFromKeytabAndReturnUGI(principal.get, keytab.get) + ugi.getShortUserName == proxyUser + } catch { + case e: IOException => + error(s"Failed to login for ${principal.get}", e) + false + } + } + } } object SparkProcessBuilder { + final val APP_KEY = "spark.app.name" + private final val CONF = "--conf" private final val CLASS = "--class" private final val PROXY_USER = "--proxy-user" - final val APP_KEY = "spark.app.name" + private final val PRINCIPAL = "spark.kerberos.principal" + private final val KEYTAB = "spark.kerberos.keytab" } diff --git a/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala b/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala index 0296b8a4728..cd54ab9fc24 100644 --- a/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala +++ b/kyuubi-main/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala @@ -20,11 +20,12 @@ package org.apache.kyuubi.engine.spark import java.nio.file.{Files, Paths} import java.util.concurrent.TimeUnit -import org.apache.kyuubi.{KyuubiFunSuite, KyuubiSQLException} +import org.apache.kyuubi.{KerberizedTestHelper, KyuubiSQLException, Utils} import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf._ +import org.apache.kyuubi.service.ServiceUtils -class SparkProcessBuilderSuite extends KyuubiFunSuite { +class SparkProcessBuilderSuite extends KerberizedTestHelper { private val conf = KyuubiConf() .set(EMBEDDED_ZK_PORT, 5555) .set(EMBEDDED_ZK_TEMP_DIR, "spark_process_test") @@ -71,4 +72,31 @@ class SparkProcessBuilderSuite extends KyuubiFunSuite { error1.getMessage.contains("Caused by: org.apache.hadoop.hive.ql.metadata.HiveException:")) } + test("proxy user or keytab") { + val b1 = new SparkProcessBuilder("kentyao", conf) + assert(b1.toString.contains("--proxy-user kentyao")) + + val conf1 = conf ++ Map("spark.kerberos.principal" -> testPrincipal) + val b2 = new SparkProcessBuilder("kentyao", conf1) + assert(b2.toString.contains("--proxy-user kentyao")) + + val conf2 = conf ++ Map("spark.kerberos.keytab" -> testKeytab) + val b3 = new SparkProcessBuilder("kentyao", conf2) + assert(b3.toString.contains("--proxy-user kentyao")) + + tryWithSecurityEnabled { + val conf3 = conf ++ Map("spark.kerberos.principal" -> testPrincipal, + "spark.kerberos.keytab" -> "testKeytab") + val b4 = new SparkProcessBuilder(Utils.currentUser, conf3) + assert(b4.toString.contains(s"--proxy-user ${Utils.currentUser}")) + + val conf4 = conf ++ Map("spark.kerberos.principal" -> testPrincipal, + "spark.kerberos.keytab" -> testKeytab) + val b5 = new SparkProcessBuilder("kentyao", conf4) + assert(b5.toString.contains("--proxy-user kentyao")) + + val b6 = new SparkProcessBuilder(ServiceUtils.getShortName(testPrincipal), conf4) + assert(!b6.toString.contains("--proxy-user kentyao")) + } + } }