get

apache · Feb 18, 2023 · 6da9f9b · 6da9f9b
1 parent ba1a868
commit 6da9f9b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 7 deletions.
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
@@ -131,12 +131,11 @@ object KyuubiServer extends Logging {
   }
 
   private[kyuubi] def refreshUnlimitedUsers(): Unit = synchronized {
-    val existingUnlimitedUsers =
-      kyuubiServer.conf.get(KyuubiConf.SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet
-    val refreshedUnlimitedUsers = KyuubiConf().loadFileDefaults().get(
-      KyuubiConf.SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet
-    kyuubiServer.backendService.sessionManager.asInstanceOf[KyuubiSessionManager]
-      .refreshUnlimitedUsers(refreshedUnlimitedUsers)
+    val sessionMgr = kyuubiServer.backendService.sessionManager.asInstanceOf[KyuubiSessionManager]
+    val existingUnlimitedUsers = sessionMgr.getUnlimitedUsers()
+    sessionMgr.refreshUnlimitedUsers(KyuubiConf().loadFileDefaults().get(
+      KyuubiConf.SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet)
+    val refreshedUnlimitedUsers = sessionMgr.getUnlimitedUsers()
     info(s"Refreshed unlimited users from $existingUnlimitedUsers to $refreshedUnlimitedUsers")
   }
 }

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
@@ -300,6 +300,10 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
       userUnlimitedList)
   }
 
+  private[kyuubi] def getUnlimitedUsers(): Set[String] = {
+    limiter.orElse(batchLimiter).map(SessionLimiter.getUnlimitedUsers).getOrElse(Set.empty)
+  }
+
   private[kyuubi] def refreshUnlimitedUsers(userUnlimitedList: Set[String]): Unit = {
     limiter.foreach(SessionLimiter.resetUnlimitedUsers(_, userUnlimitedList))
     batchLimiter.foreach(SessionLimiter.resetUnlimitedUsers(_, userUnlimitedList))

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala
@@ -138,10 +138,14 @@ object SessionLimiter {
       unlimitedUsers)
   }
 
-  def resetUnlimitedUsers(limiter: SessionLimiter, unlimitedUsers: Set[String]): Unit = {
+  def resetUnlimitedUsers(limiter: SessionLimiter, unlimitedUsers: Set[String]): Unit =
     limiter match {
       case l: SessionLimiterWithUnlimitedUsersImpl => l.setUnlimitedUsers(unlimitedUsers)
       case _ =>
     }
+
+  def getUnlimitedUsers(limiter: SessionLimiter): Set[String] = limiter match {
+    case l: SessionLimiterWithUnlimitedUsersImpl => l.unlimitedUsers
+    case _ => Set.empty
   }
 }