saev

apache · Feb 6, 2024 · c1979c7 · c1979c7
1 parent 5ef4390
commit c1979c7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/AdminRestApi.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/AdminRestApi.java
@@ -84,14 +84,17 @@ public List<Engine> listEngines(
   }
 
   public List<SessionData> listSessions() {
-    return listSessions(Collections.emptyList());
+    return listSessions(Collections.emptyList(), null);
   }
 
-  public List<SessionData> listSessions(List<String> users) {
+  public List<SessionData> listSessions(List<String> users, String sessionType) {
     Map<String, Object> params = new HashMap<>();
     if (users != null && !users.isEmpty()) {
       params.put("users", String.join(",", users));
     }
+    if (StringUtils.isNotBlank(sessionType)) {
+      params.put("sessionType", sessionType);
+    }
     SessionData[] result =
         this.getClient()
             .get(API_BASE_PATH + "/sessions", params, SessionData[].class, client.getAuthHeader());

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/AdminResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/AdminResource.scala
@@ -149,7 +149,9 @@ private[v1] class AdminResource extends ApiRequestContext with Logging {
     description = "get the list of all live sessions")
   @GET
   @Path("sessions")
-  def sessions(@QueryParam("users") users: String): Seq[SessionData] = {
+  def sessions(
+      @QueryParam("users") users: String,
+      @QueryParam("sessionType") sessionType: String): Seq[SessionData] = {
     val userName = fe.getSessionUser(Map.empty[String, String])
     val ipAddress = fe.getIpAddress
     info(s"Received listing all live sessions request from $userName/$ipAddress")
@@ -158,6 +160,10 @@ private[v1] class AdminResource extends ApiRequestContext with Logging {
         s"$userName is not allowed to list all live sessions")
     }
     var sessions = fe.be.sessionManager.allSessions()
+    if (StringUtils.isNoneBlank(sessionType)) {
+      sessions = sessions.filter(session =>
+        sessionType.equals(session.asInstanceOf[KyuubiSession].sessionType.toString))
+    }
     if (StringUtils.isNotBlank(users)) {
       val usersSet = users.split(",").toSet
       sessions = sessions.filter(session => usersSet.contains(session.user))