From e5ed088181405ce79b8667a989ca9e91d4b215c2 Mon Sep 17 00:00:00 2001
From: Kent Yao <yaooqinn@hotmail.com>
Date: Fri, 16 Nov 2018 14:02:31 +0800
Subject: [PATCH] use session ugi to collect tokens

---
 .../main/scala/yaooqinn/kyuubi/session/KyuubiSession.scala | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/kyuubi-server/src/main/scala/yaooqinn/kyuubi/session/KyuubiSession.scala b/kyuubi-server/src/main/scala/yaooqinn/kyuubi/session/KyuubiSession.scala
index 738637997..b22754c6d 100644
--- a/kyuubi-server/src/main/scala/yaooqinn/kyuubi/session/KyuubiSession.scala
+++ b/kyuubi-server/src/main/scala/yaooqinn/kyuubi/session/KyuubiSession.scala
@@ -36,6 +36,7 @@ import yaooqinn.kyuubi.operation.{KyuubiOperation, OperationHandle, OperationMan
 import yaooqinn.kyuubi.schema.RowSet
 import yaooqinn.kyuubi.session.security.TokenCollector
 import yaooqinn.kyuubi.spark.SparkSessionWithUGI
+import yaooqinn.kyuubi.utils.KyuubiHadoopUtil
 
 /**
  * An Execution Session with [[SparkSession]] instance inside, which shares [[SparkContext]]
@@ -75,7 +76,7 @@ private[kyuubi] class KyuubiSession(
           currentUser.reloginFromKeytab()
         }
         val user = UserGroupInformation.createProxyUser(username, currentUser)
-        TokenCollector.obtainTokenIfRequired(conf, user.getCredentials)
+        KyuubiHadoopUtil.doAs(user)(TokenCollector.obtainTokenIfRequired(conf, user.getCredentials))
         user
       } else {
         UserGroupInformation.createRemoteUser(username)
@@ -139,7 +140,9 @@ private[kyuubi] class KyuubiSession(
 
   def ugi: UserGroupInformation = {
     val creds = new Credentials
-    TokenCollector.obtainTokenIfRequired(conf, creds)
+    KyuubiHadoopUtil.doAs(sessionUGI) {
+      TokenCollector.obtainTokenIfRequired(conf, creds)
+    }
     sessionUGI.addCredentials(creds)
     sessionUGI
   }