From f3986d2d677b70eae9bf5515eef47e88e3748a41 Mon Sep 17 00:00:00 2001
From: zwangsheng <2213335496@qq.com>
Date: Tue, 28 Mar 2023 09:44:49 +0800
Subject: [PATCH] [KYUUBI #4618][REST] Admin Resource list operations with
 sessionHandle filter

### _Why are the changes needed?_

Close #4618
`api/v1/admin/operations` add queryParams sessionHandle

### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [ ] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #4621 from zwangsheng/KYUUBI_4618.

Closes #4618

1f1e40213 [zwangsheng] [KYUUBI #4618][REST] Admin Resource list operations with sessionhandle filter

Authored-by: zwangsheng <2213335496@qq.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../kyuubi/server/api/v1/AdminResource.scala       |  8 +++++++-
 .../kyuubi/server/api/v1/AdminResourceSuite.scala  | 14 +++++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

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 ee63243cc..0d8b31b2c 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
@@ -160,7 +160,9 @@ private[v1] class AdminResource extends ApiRequestContext with Logging {
       "get the list of all active operations")
   @GET
   @Path("operations")
-  def listOperations(@QueryParam("users") users: String): Seq[OperationData] = {
+  def listOperations(
+      @QueryParam("users") users: String,
+      @QueryParam("sessionHandle") sessionHandle: String): Seq[OperationData] = {
     val userName = fe.getSessionUser(Map.empty[String, String])
     val ipAddress = fe.getIpAddress
     info(s"Received listing all of the active operations request from $userName/$ipAddress")
@@ -173,6 +175,10 @@ private[v1] class AdminResource extends ApiRequestContext with Logging {
       val usersSet = users.split(",").toSet
       operations = operations.filter(operation => usersSet.contains(operation.getSession.user))
     }
+    if (StringUtils.isNotBlank(sessionHandle)) {
+      operations = operations.filter(operation =>
+        operation.getSession.handle.equals(SessionHandle.fromUUID(sessionHandle)))
+    }
     operations
       .map(operation => ApiUtils.operationData(operation.asInstanceOf[KyuubiOperation])).toSeq
   }
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
index e7281dc5a..a10994d7e 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala
@@ -189,7 +189,7 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
       "localhost",
       Map("testConfig" -> "testValue"))
 
-    fe.be.openSession(
+    val sessionHandle = fe.be.openSession(
       HIVE_CLI_SERVICE_PROTOCOL_V2,
       "test_user_2",
       "xxxxxx",
@@ -227,8 +227,17 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
       .request()
       .header(AUTHORIZATION_HEADER, s"BASIC $encodeAuthorization")
       .get()
-    val operations = response.readEntity(classOf[Seq[OperationData]])
+    var operations = response.readEntity(classOf[Seq[OperationData]])
     assert(operations.size == 2)
+
+    response = webTarget.path("api/v1/admin/operations")
+      .queryParam("sessionHandle", sessionHandle.identifier)
+      .request()
+      .header(AUTHORIZATION_HEADER, s"BASIC $encodeAuthorization")
+      .get()
+    operations = response.readEntity(classOf[Seq[OperationData]])
+    assert(200 == response.getStatus)
+    assert(operations.size == 1)
   }
 
   test("list/close operations") {
@@ -479,5 +488,4 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
       }
     }
   }
-
 }