[KYUUBI #5704] Add unit tests for kyuubi.session.proxy.user in REST…

…ful API # 🔍 Description ## Issue References 🔗 Close #5704 ## Describe Your Solution 🔧 Add unit tests for `kyuubi.session.proxy.user` in RESTful API. ## Types of changes 🔖 - [ ] 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) ## Test Plan 🧪 #### Behavior Without This Pull Request ⚰️ #### Behavior With This Pull Request 🎉 #### Related Unit Tests --- # Checklists ## 📝 Author Self Checklist - [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 - [x] 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) ## 📝 Committer Pre-Merge Checklist - [x] Pull request title is okay. - [x] No license issues. - [x] Milestone correctly set? - [x] Test coverage is ok - [x] Assignees are selected. - [x] Minimum number of approvals - [x] No changes are requested **Be nice. Be informative.** Closes #5705 from mrtisttt/add-test-for-kyuubi-proxy-user. Closes #5704 7da5fe3 [mrtisttt] Fix spotless 2ab8fc4 [mrtisttt] Fix scalastyle d050476 [Cheng Pan] Update kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala a955f1f [Cheng Pan] Update kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala e2f1dad [Cheng Pan] Update kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/AdminResourceSuite.scala 91d9bc5 [Cheng Pan] Update kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala 71770a6 [Cheng Pan] Update kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala 1a693c8 [mrtisttt] fix checkstyle b648a31 [mrtisttt] Delete test of delete batch with proxyUser and improve the code based on Review comments c592c24 [mrtisttt] Refactor the code to make it cleaner and more graceful. a75f1a1 [mrtisttt] Fix spotless violations cf09958 [mrtisttt] Add unit tests for `kyuubi.session.proxy.user` in RESTful API Lead-authored-by: mrtisttt <[email protected]> Co-authored-by: Cheng Pan <[email protected]> Signed-off-by: Cheng Pan <[email protected]>
apache · Nov 22, 2023 · 111786f · 111786f
1 parent 096be39
commit 111786f
Show file tree

Hide file tree

Showing 2 changed files with 168 additions and 1 deletion.
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
@@ -20,7 +20,7 @@ package org.apache.kyuubi.server.api.v1
 import java.time.Duration
 import java.util.UUID
 import javax.ws.rs.client.Entity
-import javax.ws.rs.core.{GenericType, MediaType}
+import javax.ws.rs.core.{GenericType, MediaType, Response}
 
 import scala.collection.JavaConverters._
 
@@ -387,6 +387,69 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
     }
   }
 
+  test("delete engine - user share level & proxyUser") {
+    val normalUser = "kyuubi"
+
+    val id = UUID.randomUUID().toString
+    conf.set(KyuubiConf.ENGINE_SHARE_LEVEL, USER.toString)
+    conf.set(KyuubiConf.ENGINE_TYPE, SPARK_SQL.toString)
+    conf.set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)
+    conf.set(HighAvailabilityConf.HA_NAMESPACE, "kyuubi_test")
+    conf.set(KyuubiConf.GROUP_PROVIDER, "hadoop")
+
+    // In EngineRef, when use hive.server2.proxy.user or kyuubi.session.proxy.user
+    // the user is the proxyUser, and in our test it is normalUser
+    val engine =
+      new EngineRef(conf.clone, user = normalUser, PluginLoader.loadGroupProvider(conf), id, null)
+
+    // so as the firstChild in engineSpace we use normalUser
+    val engineSpace = DiscoveryPaths.makePath(
+      s"kyuubi_test_${KYUUBI_VERSION}_USER_SPARK_SQL",
+      normalUser,
+      "default")
+
+    withDiscoveryClient(conf) { client =>
+      engine.getOrCreate(client)
+
+      assert(client.pathExists(engineSpace))
+      assert(client.getChildren(engineSpace).size == 1)
+
+      def runDeleteEngine(
+          kyuubiProxyUser: Option[String],
+          hs2ProxyUser: Option[String]): Response = {
+        var internalWebTarget = webTarget.path("api/v1/admin/engine")
+          .queryParam("sharelevel", "USER")
+          .queryParam("type", "SPARK_SQL")
+
+        kyuubiProxyUser.map(username =>
+          internalWebTarget = internalWebTarget.queryParam("proxyUser", username))
+        hs2ProxyUser.map(username =>
+          internalWebTarget = internalWebTarget.queryParam("hive.server2.proxy.user", username))
+
+        internalWebTarget.request(MediaType.APPLICATION_JSON_TYPE)
+          .header(AUTHORIZATION_HEADER, HttpAuthUtils.basicAuthorizationHeader("anonymous"))
+          .delete()
+      }
+
+      // use proxyUser
+      val deleteEngineResponse1 = runDeleteEngine(Option(normalUser), None)
+      assert(deleteEngineResponse1.getStatus === 405)
+      val errorMessage = s"Failed to validate proxy privilege of anonymous for $normalUser"
+      assert(deleteEngineResponse1.readEntity(classOf[String]).contains(errorMessage))
+
+      // it should be the same behavior as hive.server2.proxy.user
+      val deleteEngineResponse2 = runDeleteEngine(None, Option(normalUser))
+      assert(deleteEngineResponse2.getStatus === 405)
+      assert(deleteEngineResponse2.readEntity(classOf[String]).contains(errorMessage))
+
+      // when both set, proxyUser takes precedence
+      val deleteEngineResponse3 =
+        runDeleteEngine(Option(normalUser), Option(s"${normalUser}HiveServer2"))
+      assert(deleteEngineResponse3.getStatus === 405)
+      assert(deleteEngineResponse3.readEntity(classOf[String]).contains(errorMessage))
+    }
+  }
+
   test("list engine - user share level") {
     val id = UUID.randomUUID().toString
     conf.set(KyuubiConf.ENGINE_SHARE_LEVEL, USER.toString)
@@ -545,6 +608,69 @@ class AdminResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
     }
   }
 
+  test("list engine - user share level & proxyUser") {
+    val normalUser = "kyuubi"
+
+    val id = UUID.randomUUID().toString
+    conf.set(KyuubiConf.ENGINE_SHARE_LEVEL, USER.toString)
+    conf.set(KyuubiConf.ENGINE_TYPE, SPARK_SQL.toString)
+    conf.set(KyuubiConf.FRONTEND_THRIFT_BINARY_BIND_PORT, 0)
+    conf.set(HighAvailabilityConf.HA_NAMESPACE, "kyuubi_test")
+    conf.set(KyuubiConf.GROUP_PROVIDER, "hadoop")
+
+    // In EngineRef, when use hive.server2.proxy.user or kyuubi.session.proxy.user
+    // the user is the proxyUser, and in our test it is normalUser
+    val engine =
+      new EngineRef(conf.clone, user = normalUser, PluginLoader.loadGroupProvider(conf), id, null)
+
+    // so as the firstChild in engineSpace we use normalUser
+    val engineSpace = DiscoveryPaths.makePath(
+      s"kyuubi_test_${KYUUBI_VERSION}_USER_SPARK_SQL",
+      normalUser,
+      "")
+
+    withDiscoveryClient(conf) { client =>
+      engine.getOrCreate(client)
+
+      assert(client.pathExists(engineSpace))
+      assert(client.getChildren(engineSpace).size == 1)
+
+      def runListEngine(kyuubiProxyUser: Option[String], hs2ProxyUser: Option[String]): Response = {
+        var internalWebTarget = webTarget.path("api/v1/admin/engine")
+          .queryParam("sharelevel", "USER")
+          .queryParam("type", "SPARK_SQL")
+
+        kyuubiProxyUser.map { username =>
+          internalWebTarget = internalWebTarget.queryParam("proxyUser", username)
+        }
+        hs2ProxyUser.map { username =>
+          internalWebTarget = internalWebTarget.queryParam("hive.server2.proxy.user", username)
+        }
+
+        internalWebTarget.request(MediaType.APPLICATION_JSON_TYPE)
+          .header(AUTHORIZATION_HEADER, HttpAuthUtils.basicAuthorizationHeader("anonymous"))
+          .get
+      }
+
+      // use proxyUser
+      val listEngineResponse1 = runListEngine(Option(normalUser), None)
+      assert(listEngineResponse1.getStatus === 405)
+      val errorMessage = s"Failed to validate proxy privilege of anonymous for $normalUser"
+      assert(listEngineResponse1.readEntity(classOf[String]).contains(errorMessage))
+
+      // it should be the same behavior as hive.server2.proxy.user
+      val listEngineResponse2 = runListEngine(None, Option(normalUser))
+      assert(listEngineResponse2.getStatus === 405)
+      assert(listEngineResponse2.readEntity(classOf[String]).contains(errorMessage))
+
+      // when both set, proxyUser takes precedence
+      val listEngineResponse3 =
+        runListEngine(Option(normalUser), Option(s"${normalUser}HiveServer2"))
+      assert(listEngineResponse3.getStatus === 405)
+      assert(listEngineResponse3.readEntity(classOf[String]).contains(errorMessage))
+    }
+  }
+
   test("list server") {
     // Mock Kyuubi Server
     val serverDiscovery = mock[ServiceDiscovery]

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
@@ -24,6 +24,7 @@ import javax.ws.rs.client.Entity
 import javax.ws.rs.core.{MediaType, Response}
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration.DurationInt
 
@@ -842,4 +843,44 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite
     val getBatchListResponse = response.readEntity(classOf[GetBatchesResponse])
     assert(getBatchListResponse.getTotal == 1)
   }
+
+  test("open batch session with proxyUser") {
+    val normalUser = "kyuubi"
+
+    def runOpenBatchExecutor(
+        kyuubiProxyUser: Option[String],
+        hs2ProxyUser: Option[String]): Response = {
+      val conf = mutable.Map("spark.master" -> "local")
+
+      kyuubiProxyUser.map { username =>
+        conf += (PROXY_USER.key -> username)
+      }
+      hs2ProxyUser.map { username =>
+        conf += (KyuubiAuthenticationFactory.HS2_PROXY_USER -> username)
+      }
+      val proxyUserRequest = newSparkBatchRequest(conf.toMap)
+
+      webTarget.path("api/v1/batches")
+        .request(MediaType.APPLICATION_JSON_TYPE)
+        .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("anonymous"))
+        .post(Entity.entity(proxyUserRequest, MediaType.APPLICATION_JSON_TYPE))
+    }
+
+    // use kyuubi.session.proxy.user
+    val proxyUserResponse1 = runOpenBatchExecutor(Option(normalUser), None)
+    assert(proxyUserResponse1.getStatus === 405)
+    val errorMessage = s"Failed to validate proxy privilege of anonymous for $normalUser"
+    assert(proxyUserResponse1.readEntity(classOf[String]).contains(errorMessage))
+
+    // it should be the same behavior as hive.server2.proxy.user
+    val proxyUserResponse2 = runOpenBatchExecutor(None, Option(normalUser))
+    assert(proxyUserResponse2.getStatus === 405)
+    assert(proxyUserResponse2.readEntity(classOf[String]).contains(errorMessage))
+
+    // when both set, kyuubi.session.proxy.user takes precedence
+    val proxyUserResponse3 =
+      runOpenBatchExecutor(Option(normalUser), Option(s"${normalUser}HiveServer2"))
+    assert(proxyUserResponse3.getStatus === 405)
+    assert(proxyUserResponse3.readEntity(classOf[String]).contains(errorMessage))
+  }
 }