From 22d7910777a3f83f14935e85acb54f934cf3ecf3 Mon Sep 17 00:00:00 2001
From: fwang12 <fwang12@ebay.com>
Date: Fri, 17 Nov 2023 20:52:40 +0800
Subject: [PATCH] [KYUUBI #5717] Infer the proxy user automatically for delete
 batch operation

Infer the batch user from session or metadata, user do not need to specify the proxy user anymore.

This pr also align the behavior of BatchesResource with that of SessionsResource and OperationsResource(no proxy user parameter).

For Kyuubi Batch, Session and Operation, these resources have the user attiribute.

So we only need to check whether the authentication user has the permission to access the resource.

This pull request fixes #

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

- [ ] Bugfix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)

---

- [x] My code follows the [style guidelines](https://kyuubi.readthedocs.io/en/master/contributing/code/style.html) of this project
- [x] I have performed a self-review
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html)

- [x] Pull request title is okay.
- [x] No license issues.
- [x] Milestone correctly set?
- [ ] Test coverage is ok
- [x] Assignees are selected.
- [x] Minimum number of approvals
- [x] No changes are requested

**Be nice. Be informative.**

Closes #5717 from turboFei/hive_server2_proxy_user.

Closes #5717

70ad7e76d [fwang12] comment
c721a751a [fwang12] ignore
da92bd5a1 [fwang12] fix ut
9a197d005 [fwang12] doc
c8ed5f9cf [fwang12] ut
cef9e329c [fwang12] do not use proxy user

Authored-by: fwang12 <fwang12@ebay.com>
Signed-off-by: fwang12 <fwang12@ebay.com>
(cherry picked from commit 3478fc9dfb4716dd4a808506f3dab26155eaba91)
Signed-off-by: fwang12 <fwang12@ebay.com>
---
 docs/client/rest/rest_api.md                  |  6 -----
 docs/deployment/migration-guide.md            |  4 +++
 .../ctl/cmd/delete/DeleteBatchCommand.scala   |  2 +-
 .../kyuubi/ctl/BatchCliArgumentsSuite.scala   | 12 ---------
 .../apache/kyuubi/client/BatchRestApi.java    | 16 +++++++++++
 .../kyuubi/client/BatchRestClientTest.java    |  4 +--
 .../server/api/v1/BatchesResource.scala       | 27 +++++--------------
 .../server/api/v1/InternalRestClient.scala    |  2 +-
 .../server/api/v1/BatchesResourceSuite.scala  | 12 +--------
 .../rest/client/BatchRestApiSuite.scala       | 11 ++------
 10 files changed, 34 insertions(+), 62 deletions(-)

diff --git a/docs/client/rest/rest_api.md b/docs/client/rest/rest_api.md
index 2fa2a7bf703..c008202f81b 100644
--- a/docs/client/rest/rest_api.md
+++ b/docs/client/rest/rest_api.md
@@ -410,12 +410,6 @@ The [Batch](#batch).
 
 Kill the batch if it is still running.
 
-#### Request Parameters
-
-| Name                    | Description                   | Type             |
-|:------------------------|:------------------------------|:-----------------|
-| hive.server2.proxy.user | the proxy user to impersonate | String(optional) |
-
 #### Response Body
 
 | Name    | Description                           | Type    |
diff --git a/docs/deployment/migration-guide.md b/docs/deployment/migration-guide.md
index 0430e8cff8d..b8d9798bfcf 100644
--- a/docs/deployment/migration-guide.md
+++ b/docs/deployment/migration-guide.md
@@ -17,6 +17,10 @@
 
 # Kyuubi Migration Guide
 
+## Upgrading from Kyuubi 1.8.0 to 1.8.1
+
+* Since Kyuubi 1.8.1, for `DELETE /batches/${batchId}`, `hive.server2.proxy.user` is not needed in the request parameters.
+
 ## Upgrading from Kyuubi 1.7 to 1.8
 
 * Since Kyuubi 1.8, SQLite is added and becomes the default database type of Kyuubi metastore, as Derby has been deprecated.
diff --git a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/cmd/delete/DeleteBatchCommand.scala b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/cmd/delete/DeleteBatchCommand.scala
index 3988620adb8..ee4a14d2666 100644
--- a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/cmd/delete/DeleteBatchCommand.scala
+++ b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/cmd/delete/DeleteBatchCommand.scala
@@ -36,7 +36,7 @@ class DeleteBatchCommand(cliConfig: CliConfig) extends Command[Batch](cliConfig)
       val batchRestApi: BatchRestApi = new BatchRestApi(kyuubiRestClient)
       val batchId = normalizedCliConfig.batchOpts.batchId
 
-      val result = batchRestApi.deleteBatch(batchId, normalizedCliConfig.commonOpts.hs2ProxyUser)
+      val result = batchRestApi.deleteBatch(batchId)
 
       info(JsonUtils.toJson(result))
 
diff --git a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/BatchCliArgumentsSuite.scala b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/BatchCliArgumentsSuite.scala
index bf8f101e00a..5987ac16338 100644
--- a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/BatchCliArgumentsSuite.scala
+++ b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/BatchCliArgumentsSuite.scala
@@ -119,18 +119,6 @@ class BatchCliArgumentsSuite extends KyuubiFunSuite with TestPrematureExit {
     }
   }
 
-  test("delete batch with hs2ProxyUser") {
-    val args = Array(
-      "delete",
-      "batch",
-      "f7fd702c-e54e-11ec-8fea-0242ac120002",
-      "--hs2ProxyUser",
-      "b_user")
-    val opArgs = new ControlCliArguments(args)
-    assert(opArgs.cliConfig.batchOpts.batchId == "f7fd702c-e54e-11ec-8fea-0242ac120002")
-    assert(opArgs.cliConfig.commonOpts.hs2ProxyUser == "b_user")
-  }
-
   test("test list batch option") {
     val args = Array(
       "list",
diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/BatchRestApi.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/BatchRestApi.java
index 7d113308df1..e6f9577b345 100644
--- a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/BatchRestApi.java
+++ b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/BatchRestApi.java
@@ -23,8 +23,11 @@
 import org.apache.kyuubi.client.api.v1.dto.*;
 import org.apache.kyuubi.client.util.JsonUtils;
 import org.apache.kyuubi.client.util.VersionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class BatchRestApi {
+  static final Logger LOG = LoggerFactory.getLogger(BatchRestApi.class);
 
   private KyuubiRestClient client;
 
@@ -101,7 +104,15 @@ public OperationLog getBatchLocalLog(String batchId, int from, int size) {
     return this.getClient().get(path, params, OperationLog.class, client.getAuthHeader());
   }
 
+  /**
+   * hs2ProxyUser for delete batch is deprecated since 1.8.1, please use {@link
+   * #deleteBatch(String)} instead.
+   */
+  @Deprecated
   public CloseBatchResponse deleteBatch(String batchId, String hs2ProxyUser) {
+    LOG.warn(
+        "The method `deleteBatch(batchId, hs2ProxyUser)` is deprecated since 1.8.1, "
+            + "using `deleteBatch(batchId)` instead.");
     Map<String, Object> params = new HashMap<>();
     params.put("hive.server2.proxy.user", hs2ProxyUser);
 
@@ -109,6 +120,11 @@ public CloseBatchResponse deleteBatch(String batchId, String hs2ProxyUser) {
     return this.getClient().delete(path, params, CloseBatchResponse.class, client.getAuthHeader());
   }
 
+  public CloseBatchResponse deleteBatch(String batchId) {
+    String path = String.format("%s/%s", API_BASE_PATH, batchId);
+    return this.getClient().delete(path, null, CloseBatchResponse.class, client.getAuthHeader());
+  }
+
   private IRestClient getClient() {
     return this.client.getHttpClient();
   }
diff --git a/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/BatchRestClientTest.java b/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/BatchRestClientTest.java
index 80fb1c4b95f..9715460ca5f 100644
--- a/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/BatchRestClientTest.java
+++ b/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/BatchRestClientTest.java
@@ -267,13 +267,13 @@ public void getOperationLogTest() {
   public void deleteBatchTest() {
     // test spnego auth
     BatchTestServlet.setAuthSchema(NEGOTIATE_AUTH);
-    CloseBatchResponse response = spnegoBatchRestApi.deleteBatch("71535", "b_test");
+    CloseBatchResponse response = spnegoBatchRestApi.deleteBatch("71535");
     assertTrue(response.isSuccess());
 
     // test basic auth
     BatchTestServlet.setAuthSchema(BASIC_AUTH);
     BatchTestServlet.allowAnonymous(false);
-    response = basicBatchRestApi.deleteBatch("71535", "b_test");
+    response = basicBatchRestApi.deleteBatch("71535");
     assertTrue(response.isSuccess());
   }
 }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
index 5e32016bf14..58e70319032 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
@@ -23,7 +23,6 @@ import java.util.{Collections, Locale, UUID}
 import java.util.concurrent.ConcurrentHashMap
 import javax.ws.rs._
 import javax.ws.rs.core.MediaType
-import javax.ws.rs.core.Response.Status
 
 import scala.collection.JavaConverters._
 import scala.util.{Failure, Success, Try}
@@ -448,18 +447,7 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
     description = "close and cancel a batch session")
   @DELETE
   @Path("{batchId}")
-  def closeBatchSession(
-      @PathParam("batchId") batchId: String,
-      @QueryParam("hive.server2.proxy.user") hs2ProxyUser: String): CloseBatchResponse = {
-
-    def checkPermission(operator: String, owner: String): Unit = {
-      if (operator != owner) {
-        throw new WebApplicationException(
-          s"$operator is not allowed to close the session belong to $owner",
-          Status.METHOD_NOT_ALLOWED)
-      }
-    }
-
+  def closeBatchSession(@PathParam("batchId") batchId: String): CloseBatchResponse = {
     def forceKill(
         appMgrInfo: ApplicationManagerInfo,
         batchId: String,
@@ -472,16 +460,15 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
     }
 
     val sessionHandle = formatSessionHandle(batchId)
-    val userName = fe.getSessionUser(hs2ProxyUser)
 
     sessionManager.getBatchSession(sessionHandle).map { batchSession =>
-      checkPermission(userName, batchSession.user)
+      fe.getSessionUser(batchSession.user)
       sessionManager.closeSession(batchSession.handle)
       val (killed, msg) = batchSession.batchJobSubmissionOp.getKillMessage
       new CloseBatchResponse(killed, msg)
     }.getOrElse {
       sessionManager.getBatchMetadata(batchId).map { metadata =>
-        checkPermission(userName, metadata.username)
+        fe.getSessionUser(metadata.username)
         if (OperationState.isTerminal(OperationState.withName(metadata.state))) {
           new CloseBatchResponse(false, s"The batch[$metadata] has been terminated.")
         } else if (batchV2Enabled(metadata.requestConf) && metadata.state == "INITIALIZED" &&
@@ -492,21 +479,21 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
         } else if (batchV2Enabled(metadata.requestConf) && metadata.kyuubiInstance == null) {
           // code goes here indicates metadata is outdated, recursively calls itself to refresh
           // the metadata
-          closeBatchSession(batchId, hs2ProxyUser)
+          closeBatchSession(batchId)
         } else if (metadata.kyuubiInstance != fe.connectionUrl) {
           info(s"Redirecting delete batch[$batchId] to ${metadata.kyuubiInstance}")
           val internalRestClient = getInternalRestClient(metadata.kyuubiInstance)
           try {
-            internalRestClient.deleteBatch(userName, batchId)
+            internalRestClient.deleteBatch(metadata.username, batchId)
           } catch {
             case e: KyuubiRestException =>
               error(s"Error redirecting delete batch[$batchId] to ${metadata.kyuubiInstance}", e)
-              val (killed, msg) = forceKill(metadata.appMgrInfo, batchId, userName)
+              val (killed, msg) = forceKill(metadata.appMgrInfo, batchId, metadata.username)
               new CloseBatchResponse(killed, if (killed) msg else Utils.stringifyException(e))
           }
         } else { // should not happen, but handle this for safe
           warn(s"Something wrong on deleting batch[$batchId], try forcibly killing application")
-          val (killed, msg) = forceKill(metadata.appMgrInfo, batchId, userName)
+          val (killed, msg) = forceKill(metadata.appMgrInfo, batchId, metadata.username)
           new CloseBatchResponse(killed, msg)
         }
       }.getOrElse {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
index dbabc5882a5..59d14dacd1e 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
@@ -60,7 +60,7 @@ class InternalRestClient(
 
   def deleteBatch(user: String, batchId: String): CloseBatchResponse = {
     withAuthUser(user) {
-      internalBatchRestApi.deleteBatch(batchId, null)
+      internalBatchRestApi.deleteBatch(batchId)
     }
   }
 
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
index fda47d58351..0d836a05aaf 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala
@@ -211,7 +211,7 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite
       .header(AUTHORIZATION_HEADER, basicAuthorizationHeader(batch.getId))
       .delete()
     assert(deleteBatchResponse.getStatus === 405)
-    errorMessage = s"${batch.getId} is not allowed to close the session belong to anonymous"
+    errorMessage = s"Failed to validate proxy privilege of ${batch.getId} for anonymous"
     assert(deleteBatchResponse.readEntity(classOf[String]).contains(errorMessage))
 
     // invalid batchId
@@ -228,16 +228,6 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite
       .delete()
     assert(deleteBatchResponse.getStatus === 404)
 
-    // invalid proxy user
-    deleteBatchResponse = webTarget.path(s"api/v1/batches/${batch.getId}")
-      .queryParam("hive.server2.proxy.user", "invalidProxy")
-      .request(MediaType.APPLICATION_JSON_TYPE)
-      .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("anonymous"))
-      .delete()
-    assert(deleteBatchResponse.getStatus === 405)
-    errorMessage = "Failed to validate proxy privilege of anonymous for invalidProxy"
-    assert(deleteBatchResponse.readEntity(classOf[String]).contains(errorMessage))
-
     // check close batch session
     deleteBatchResponse = webTarget.path(s"api/v1/batches/${batch.getId}")
       .request(MediaType.APPLICATION_JSON_TYPE)
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
index d04826a9d20..20ec2fc0a5f 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
@@ -74,16 +74,9 @@ class BatchRestApiSuite extends RestClientTestHelper with BatchTestHelper {
     }
 
     // delete batch
-    val closeResp = batchRestApi.deleteBatch(batch.getId(), null)
+    val closeResp = batchRestApi.deleteBatch(batch.getId())
     assert(closeResp.getMsg.nonEmpty)
 
-    // delete batch - error
-    val e = intercept[KyuubiRestException] {
-      batchRestApi.deleteBatch(batch.getId(), "fake")
-    }
-    assert(e.getCause.toString.contains(
-      s"Failed to validate proxy privilege of ${ldapUser} for fake"))
-
     basicKyuubiRestClient.close()
   }
 
@@ -170,7 +163,7 @@ class BatchRestApiSuite extends RestClientTestHelper with BatchTestHelper {
     }
 
     // delete batch
-    val closeResp = batchRestApi.deleteBatch(batch.getId(), proxyUser)
+    val closeResp = batchRestApi.deleteBatch(batch.getId())
     assert(closeResp.getMsg.nonEmpty)
 
     // list batches