From c1979c7eacc70eeba729ae83e0cd4dda3738d1ed Mon Sep 17 00:00:00 2001
From: Fei Wang <fwang12@ebay.com>
Date: Mon, 5 Feb 2024 19:10:00 -0800
Subject: [PATCH] saev
---
.../main/java/org/apache/kyuubi/client/AdminRestApi.java | 7 +++++--
.../org/apache/kyuubi/server/api/v1/AdminResource.scala | 8 +++++++-
2 files changed, 12 insertions(+), 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
index 834b508de25..ef39f9c0d75 100644
--- 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
index ca35a1b6b12..8866a3443aa 100644
--- 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))