[KYUUBI #6358] Fix the incorrect construction of yarnConf

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/ApplicationOperation.scala

@@ -56,6 +56,7 @@ trait ApplicationOperation {
    * @note For implementations, please suppress exceptions and always return KillResponse
    */
   def killApplicationByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None): KillResponse
@@ -73,6 +74,7 @@ trait ApplicationOperation {
    * @return [[ApplicationInfo]]
    */
   def getApplicationInfoByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None,

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala

@@ -251,7 +251,11 @@ private[kyuubi] class EngineRef(
 
         if (started + timeout <= System.currentTimeMillis()) {
           val killMessage =
-            engineManager.killApplication(builder.appMgrInfo(), engineRefId, Some(appUser))
+            engineManager.killApplication(
+              Some(conf),
+              builder.appMgrInfo(),
+              engineRefId,
+              Some(appUser))
           builder.close(true)
           MetricsSystem.tracing(_.incCount(MetricRegistry.name(ENGINE_TIMEOUT, appUser)))
           throw KyuubiSQLException(
@@ -270,6 +274,7 @@ private[kyuubi] class EngineRef(
             }
 
             val applicationInfo = engineMgr.getApplicationInfo(
+              Some(conf),
               builder.appMgrInfo(),
               engineRefId,
               Some(appUser),
@@ -366,7 +371,7 @@ private[kyuubi] class EngineRef(
       try {
         val appMgrInfo = builder.appMgrInfo()
         builder.close(true)
-        engineManager.killApplication(appMgrInfo, engineRefId, Some(appUser))
+        engineManager.killApplication(Some(conf), appMgrInfo, engineRefId, Some(appUser))
       } catch {
         case e: Exception =>
           warn(s"Error closing engine builder, engineRefId: $engineRefId", e)

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/JpsApplicationOperation.scala

@@ -82,13 +82,15 @@ class JpsApplicationOperation extends ApplicationOperation {
   }
 
   override def killApplicationByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None): KillResponse = {
     killJpsApplicationByTag(tag, true)
   }
 
   override def getApplicationInfoByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None,

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala

@@ -184,6 +184,7 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
   }
 
   override def killApplicationByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None): KillResponse = {
@@ -224,6 +225,7 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
   }
 
   override def getApplicationInfoByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None,

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KyuubiApplicationManager.scala

@@ -61,13 +61,14 @@ class KyuubiApplicationManager extends AbstractService("KyuubiApplicationManager
   }
 
   def killApplication(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None): KillResponse = {
     var (killed, lastMessage): KillResponse = (false, null)
     for (operation <- operations if !killed) {
       if (operation.isSupported(appMgrInfo)) {
-        val (k, m) = operation.killApplicationByTag(appMgrInfo, tag, proxyUser)
+        val (k, m) = operation.killApplicationByTag(sessionConf, appMgrInfo, tag, proxyUser)
         killed = k
         lastMessage = m
       }
@@ -84,13 +85,15 @@ class KyuubiApplicationManager extends AbstractService("KyuubiApplicationManager
   }
 
   def getApplicationInfo(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None,
       submitTime: Option[Long] = None): Option[ApplicationInfo] = {
     val operation = operations.find(_.isSupported(appMgrInfo))
     operation match {
-      case Some(op) => Some(op.getApplicationInfoByTag(appMgrInfo, tag, proxyUser, submitTime))
+      case Some(op) =>
+        Some(op.getApplicationInfoByTag(sessionConf, appMgrInfo, tag, proxyUser, submitTime))
       case None => None
     }
   }

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/YarnApplicationOperation.scala

@@ -95,74 +95,139 @@ class YarnApplicationOperation extends ApplicationOperation with Logging {
     }
   }
 
+  private def withNewYarnClient[T](
+      proxyUser: Option[String],
+      sessionConf: Option[KyuubiConf])(action: YarnClient => T): T = {
+    (sessionConf) match {
+      case Some(sessionConf) =>
+        (proxyUser) match {
+          case Some(user) =>
+            Utils.doAs(user) { () =>
+              var yarnClient: YarnClient = null
+              try {
+                val yarnConf = KyuubiHadoopUtils.newYarnConfiguration(sessionConf)
+                yarnClient = createYarnClient(yarnConf)
+                action(yarnClient)
+              } finally {
+                Utils.tryLogNonFatalError(yarnClient.close())
+              }
+            }
+          case None =>
+            var yarnClient: YarnClient = null
+            try {
+              yarnClient = createAdminYarnClient(sessionConf).get
+              action(yarnClient)
+            } finally {
+              Utils.tryLogNonFatalError(yarnClient.close())
+            }
+        }
+      case None =>
+        withYarnClient(proxyUser)(action)
+    }
+  }
+
+  def createAdminYarnClient(conf: KyuubiConf): Option[YarnClient] = {
+
+    val yarnConf = KyuubiHadoopUtils.newYarnConfiguration(conf)
+
+    def createYarnClientWithCurrentUser(): Option[YarnClient] = {
+      val c = createYarnClient(yarnConf)
+      info(s"Creating admin YARN client with current user: ${Utils.currentUser}.")
+      Some(c);
+    }
+
+    def createYarnClientWithProxyUser(proxyUser: String): Option[YarnClient] =
+      Utils.doAs(proxyUser) { () =>
+        val c = createYarnClient(yarnConf)
+        info(s"Creating admin YARN client with proxy user: $proxyUser.")
+        Some(c);
+      }
+
+    YarnUserStrategy.withName(conf.get(KyuubiConf.YARN_USER_STRATEGY)) match {
+      case NONE =>
+        createYarnClientWithCurrentUser()
+      case ADMIN if conf.get(KyuubiConf.YARN_USER_ADMIN) == Utils.currentUser =>
+        createYarnClientWithCurrentUser()
+      case ADMIN =>
+        createYarnClientWithProxyUser(conf.get(KyuubiConf.YARN_USER_ADMIN))
+      case OWNER =>
+        info("Skip initializing admin YARN client")
+        None
+    }
+  }
+
   override def isSupported(appMgrInfo: ApplicationManagerInfo): Boolean =
     appMgrInfo.resourceManager.exists(_.toLowerCase(Locale.ROOT).startsWith("yarn"))
 
   override def killApplicationByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
-      proxyUser: Option[String] = None): KillResponse = withYarnClient(proxyUser) { yarnClient =>
-    try {
-      val reports = yarnClient.getApplications(null, null, Set(tag).asJava)
-      if (reports.isEmpty) {
-        (false, NOT_FOUND)
-      } else {
-        try {
-          val applicationId = reports.get(0).getApplicationId
-          yarnClient.killApplication(applicationId)
-          (true, s"Succeeded to terminate: $applicationId with $tag")
-        } catch {
-          case e: Exception =>
-            (false, s"Failed to terminate application with $tag, due to ${e.getMessage}")
+      proxyUser: Option[String] = None): KillResponse =
+    withNewYarnClient(proxyUser, sessionConf) { yarnClient =>
+      try {
+        val reports = yarnClient.getApplications(null, null, Set(tag).asJava)
+        if (reports.isEmpty) {
+          (false, NOT_FOUND)
+        } else {
+          try {
+            val applicationId = reports.get(0).getApplicationId
+            yarnClient.killApplication(applicationId)
+            (true, s"Succeeded to terminate: $applicationId with $tag")
+          } catch {
+            case e: Exception =>
+              (false, s"Failed to terminate application with $tag, due to ${e.getMessage}")
+          }
         }
+      } catch {
+        case e: Exception =>
+          (
+            false,
+            s"Failed to get while terminating application with tag $tag, due to ${e.getMessage}")
       }
-    } catch {
-      case e: Exception =>
-        (
-          false,
-          s"Failed to get while terminating application with tag $tag, due to ${e.getMessage}")
     }
-  }
 
   override def getApplicationInfoByTag(
+      sessionConf: Option[KyuubiConf],
       appMgrInfo: ApplicationManagerInfo,
       tag: String,
       proxyUser: Option[String] = None,
-      submitTime: Option[Long] = None): ApplicationInfo = withYarnClient(proxyUser) { yarnClient =>
-    debug(s"Getting application info from YARN cluster by tag: $tag")
-    val reports = yarnClient.getApplications(null, null, Set(tag).asJava)
-    if (reports.isEmpty) {
-      debug(s"Can't find target application from YARN cluster by tag: $tag")
-      submitTime match {
-        case Some(_submitTime) =>
-          val elapsedTime = System.currentTimeMillis - _submitTime
-          if (elapsedTime < submitTimeout) {
-            info(s"Wait for YARN application[tag: $tag] to be submitted, " +
-              s"elapsed time: ${elapsedTime}ms, return ${ApplicationInfo.UNKNOWN} status")
-            ApplicationInfo.UNKNOWN
-          } else {
-            error(s"Can't find target application from YARN cluster by tag: $tag, " +
-              s"elapsed time: ${elapsedTime}ms exceeds ${ENGINE_YARN_SUBMIT_TIMEOUT.key}: " +
-              s"${submitTimeout}ms, return ${ApplicationInfo.NOT_FOUND} status")
-            ApplicationInfo.NOT_FOUND
-          }
-        case _ => ApplicationInfo.NOT_FOUND
+      submitTime: Option[Long] = None): ApplicationInfo =
+    withNewYarnClient(proxyUser, sessionConf) { yarnClient =>
+      debug(s"Getting application info from YARN cluster by tag: $tag")
+      val reports = yarnClient.getApplications(null, null, Set(tag).asJava)
+      if (reports.isEmpty) {
+        debug(s"Can't find target application from YARN cluster by tag: $tag")
+        submitTime match {
+          case Some(_submitTime) =>
+            val elapsedTime = System.currentTimeMillis - _submitTime
+            if (elapsedTime < submitTimeout) {
+              info(s"Wait for YARN application[tag: $tag] to be submitted, " +
+                s"elapsed time: ${elapsedTime}ms, return ${ApplicationInfo.UNKNOWN} status")
+              ApplicationInfo.UNKNOWN
+            } else {
+              error(s"Can't find target application from YARN cluster by tag: $tag, " +
+                s"elapsed time: ${elapsedTime}ms exceeds ${ENGINE_YARN_SUBMIT_TIMEOUT.key}: " +
+                s"${submitTimeout}ms, return ${ApplicationInfo.NOT_FOUND} status")
+              ApplicationInfo.NOT_FOUND
+            }
+          case _ => ApplicationInfo.NOT_FOUND
+        }
+      } else {
+        val report = reports.get(0)
+        val info = ApplicationInfo(
+          id = report.getApplicationId.toString,
+          name = report.getName,
+          state = toApplicationState(
+            report.getApplicationId.toString,
+            report.getYarnApplicationState,
+            report.getFinalApplicationStatus),
+          url = Option(report.getTrackingUrl),
+          error = Option(report.getDiagnostics))
+        debug(s"Successfully got application info by tag: $tag. $info")
+        info
       }
-    } else {
-      val report = reports.get(0)
-      val info = ApplicationInfo(
-        id = report.getApplicationId.toString,
-        name = report.getName,
-        state = toApplicationState(
-          report.getApplicationId.toString,
-          report.getYarnApplicationState,
-          report.getFinalApplicationStatus),
-        url = Option(report.getTrackingUrl),
-        error = Option(report.getDiagnostics))
-      debug(s"Successfully got application info by tag: $tag. $info")
-      info
     }
-  }
 
   override def stop(): Unit = adminYarnClient.foreach { yarnClient =>
     Utils.tryLogNonFatalError(yarnClient.stop())

kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala

@@ -109,6 +109,7 @@ class BatchJobSubmission(
     }
     val applicationInfo =
       applicationManager.getApplicationInfo(
+        None,
         builder.appMgrInfo(),
         batchId,
         Some(session.user),
@@ -126,7 +127,7 @@ class BatchJobSubmission(
   }
 
   private[kyuubi] def killBatchApplication(): KillResponse = {
-    applicationManager.killApplication(builder.appMgrInfo(), batchId, Some(session.user))
+    applicationManager.killApplication(None, builder.appMgrInfo(), batchId, Some(session.user))
   }
 
   private val applicationCheckInterval =

kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala

@@ -340,6 +340,7 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
             case e: KyuubiRestException =>
               error(s"Error redirecting get batch[$batchId] to ${metadata.kyuubiInstance}", e)
               val batchAppStatus = sessionManager.applicationManager.getApplicationInfo(
+                None,
                 metadata.appMgrInfo,
                 batchId,
                 Some(userName),
@@ -477,7 +478,7 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
         batchId: String,
         user: String): KillResponse = {
       val (killed, message) = sessionManager.applicationManager
-        .killApplication(appMgrInfo, batchId, Some(user))
+        .killApplication(None, appMgrInfo, batchId, Some(user))
       info(s"Mark batch[$batchId] closed by ${fe.connectionUrl}")
       sessionManager.updateMetadata(Metadata(identifier = batchId, peerInstanceClosed = true))
       (killed, message)

kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala

@@ -237,9 +237,13 @@ class KyuubiOperationPerConnectionSuite extends WithKyuubiServer with HiveJDBCTe
         }
         val engineId = sessionManager.allSessions().head.handle.identifier.toString
         // kill the engine application and wait the engine terminate
-        sessionManager.applicationManager.killApplication(ApplicationManagerInfo(None), engineId)
+        sessionManager.applicationManager.killApplication(
+          None,
+          ApplicationManagerInfo(None),
+          engineId)
         eventually(timeout(30.seconds), interval(100.milliseconds)) {
           assert(sessionManager.applicationManager.getApplicationInfo(
+            None,
             ApplicationManagerInfo(None),
             engineId)
             .exists(_.state == ApplicationState.NOT_FOUND))