[KYUUBI #3300] [FOLLOWUP] simply reflection method callings in Authz'…

…s AccessRequest ### _Why are the changes needed?_ To simplify the method calling via reflection in Authz's `AccessRequest`, including codes from pull #3308 introduced. ### _How was this patch tested?_ - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [x] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request Closes #3945 from bowenliang123/3300-followup. Closes #3300 3e3f2bf [liangbowen] move exception throwing in `invoke` method ff1c9a0 [liangbowen] import java.util.{HashMap => JHashMap} 2848083 [liangbowen] import java.util.{Set => JSet} a9b4240 [liangbowen] import java.util.{Set => JSet} 408977e [liangbowen] correct error handling in AccessRequest c785a25 [liangbowen] use invokeAs 9f9c6a1 [liangbowen] rename invokeWithCast to invokeAs 258be28 [liangbowen] introduce `invokeWithCast` in AuthZUtils for simplifying object casting f990209 [liangbowen] simplify reflection method calling in AccessRequest da767db [liangbowen] lint 3b82c31 [liangbowen] unused imports d4a71da [liangbowen] unused imports 0c57d0d [liangbowen] simply getUserGroupsFromUserStore with reflection methods Authored-by: liangbowen <liangbowen@gf.com.cn> Signed-off-by: Kent Yao <yao@apache.org>
apache · Dec 12, 2022 · 8b8c7d0 · 8b8c7d0
1 parent 70c0451
commit 8b8c7d0
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 50 deletions.
diff --git a/...i-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala b/...i-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala
@@ -169,8 +169,7 @@ object PrivilegesBuilder {
         }
 
       case u if u.nodeName == "UnresolvedRelation" =>
-        val tableNameM = u.getClass.getMethod("tableName")
-        val parts = tableNameM.invoke(u).asInstanceOf[String].split("\\.")
+        val parts = invokeAs[String](u, "tableName").split("\\.")
         val db = quote(parts.init)
         privilegeObjects += tablePrivileges(TableIdentifier(parts.last, Some(db)))
 

diff --git a/...park-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/AccessRequest.scala b/...park-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/AccessRequest.scala
@@ -17,17 +17,17 @@
 
 package org.apache.kyuubi.plugin.spark.authz.ranger
 
-import java.util
+import java.util.{HashMap => JHashMap, Set => JSet}
 import java.util.Date
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable
 
 import org.apache.hadoop.security.UserGroupInformation
 import org.apache.ranger.plugin.policyengine.{RangerAccessRequestImpl, RangerPolicyEngine}
 
 import org.apache.kyuubi.plugin.spark.authz.OperationType.OperationType
 import org.apache.kyuubi.plugin.spark.authz.ranger.AccessType._
+import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils.{invoke, invokeAs}
 
 case class AccessRequest private (accessType: AccessType) extends RangerAccessRequestImpl
 
@@ -45,66 +45,47 @@ object AccessRequest {
     req.setUserGroups(userGroups)
     req.setAction(opType.toString)
     try {
-      val getRoles = SparkRangerAdminPlugin.getClass.getMethod(
+      val roles = invokeAs[JSet[String]](
+        SparkRangerAdminPlugin,
         "getRolesFromUserAndGroups",
-        classOf[String],
-        classOf[java.util.Set[String]])
-      getRoles.setAccessible(true)
-      val roles = getRoles.invoke(SparkRangerAdminPlugin, userName, userGroups)
-      val setRoles = req.getClass.getMethod("setUserRoles", classOf[java.util.Set[String]])
-      setRoles.setAccessible(true)
-      setRoles.invoke(req, roles)
+        (classOf[String], userName),
+        (classOf[JSet[String]], userGroups))
+      invoke(req, "setUserRoles", (classOf[JSet[String]], roles))
     } catch {
-      case _: NoSuchMethodException =>
+      case _: Exception =>
     }
     req.setAccessTime(new Date())
     accessType match {
       case USE => req.setAccessType(RangerPolicyEngine.ANY_ACCESS)
       case _ => req.setAccessType(accessType.toString.toLowerCase)
     }
     try {
-      val getClusterName = SparkRangerAdminPlugin.getClass.getMethod("getClusterName")
-      getClusterName.setAccessible(true)
-      val clusterName = getClusterName.invoke(SparkRangerAdminPlugin)
-      val setClusterName = req.getClass.getMethod("setClusterName", classOf[String])
-      setClusterName.setAccessible(true)
-      setClusterName.invoke(req, clusterName)
+      val clusterName = invokeAs[String](SparkRangerAdminPlugin, "getClusterName")
+      invoke(req, "setClusterName", (classOf[String], clusterName))
     } catch {
-      case _: NoSuchMethodException =>
+      case _: Exception =>
     }
     req
   }
 
-  private def getUserGroupsFromUgi(user: UserGroupInformation): util.Set[String] = {
+  private def getUserGroupsFromUgi(user: UserGroupInformation): JSet[String] = {
     user.getGroupNames.toSet.asJava
   }
 
-  private def getUserGroupsFromUserStore(user: UserGroupInformation): Option[util.Set[String]] = {
+  private def getUserGroupsFromUserStore(user: UserGroupInformation): Option[JSet[String]] = {
     try {
-      val getUserStoreEnricher = SparkRangerAdminPlugin.getClass.getMethod(
-        "getUserStoreEnricher")
-      getUserStoreEnricher.setAccessible(true)
-      val storeEnricher = getUserStoreEnricher.invoke(SparkRangerAdminPlugin)
-
-      val getRangerUserStore = storeEnricher.getClass.getMethod("getRangerUserStore")
-      getRangerUserStore.setAccessible(true)
-      val userStore = getRangerUserStore.invoke(storeEnricher)
-
-      val getUserGroupMapping = userStore.getClass.getMethod("getUserGroupMapping")
-      getUserGroupMapping.setAccessible(true)
-
-      val userGroupMappingMap: mutable.Map[String, util.Set[String]] =
-        mapAsScalaMap(getUserGroupMapping.invoke(userStore)
-          .asInstanceOf[util.HashMap[String, util.Set[String]]])
-
-      userGroupMappingMap.get(user.getShortUserName)
+      val storeEnricher = invoke(SparkRangerAdminPlugin, "getUserStoreEnricher")
+      val userStore = invoke(storeEnricher, "getRangerUserStore")
+      val userGroupMapping =
+        invokeAs[JHashMap[String, JSet[String]]](userStore, "getUserGroupMapping")
+      Some(userGroupMapping.get(user.getShortUserName))
     } catch {
       case _: NoSuchMethodException =>
         None
     }
   }
 
-  private def getUserGroups(user: UserGroupInformation): util.Set[String] = {
+  private def getUserGroups(user: UserGroupInformation): JSet[String] = {
     if (SparkRangerAdminPlugin.useUserGroupsFromUserStoreEnabled) {
       getUserGroupsFromUserStore(user)
         .getOrElse(getUserGroupsFromUgi(user))

diff --git a/...ark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/serde/tableExtractors.scala b/...ark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/serde/tableExtractors.scala
@@ -51,7 +51,7 @@ object TableExtractor {
   def getOwner(v: AnyRef): Option[String] = {
     // org.apache.spark.sql.connector.catalog.Table
     val table = invoke(v, "table")
-    val properties = invoke(table, "properties").asInstanceOf[JMap[String, String]].asScala
+    val properties = invokeAs[JMap[String, String]](table, "properties").asScala
     properties.get("owner")
   }
 }
@@ -102,8 +102,8 @@ class ResolvedTableTableExtractor extends TableExtractor {
  */
 class IdentifierTableExtractor extends TableExtractor {
   override def apply(spark: SparkSession, v1: AnyRef): Option[Table] = {
-    val namespace = invoke(v1, "namespace").asInstanceOf[Array[String]]
-    val table = invoke(v1, "name").asInstanceOf[String]
+    val namespace = invokeAs[Array[String]](v1, "namespace")
+    val table = invokeAs[String](v1, "name")
     Some(Table(Some(quote(namespace)), table, None))
   }
 }
@@ -118,7 +118,7 @@ class DataSourceV2RelationTableExtractor extends TableExtractor {
     if (v2Relation.isEmpty) {
       None
     } else {
-      val maybeIdentifier = invoke(v2Relation.get, "identifier").asInstanceOf[Option[AnyRef]]
+      val maybeIdentifier = invokeAs[Option[AnyRef]](v2Relation.get, "identifier")
       maybeIdentifier.flatMap { id =>
         val maybeTable = new IdentifierTableExtractor().apply(spark, id)
         val maybeOwner = TableExtractor.getOwner(v2Relation.get)
@@ -133,7 +133,7 @@ class DataSourceV2RelationTableExtractor extends TableExtractor {
  */
 class LogicalRelationTableExtractor extends TableExtractor {
   override def apply(spark: SparkSession, v1: AnyRef): Option[Table] = {
-    val maybeCatalogTable = invoke(v1, "catalogTable").asInstanceOf[Option[AnyRef]]
+    val maybeCatalogTable = invokeAs[Option[AnyRef]](v1, "catalogTable")
     maybeCatalogTable.flatMap { ct =>
       new CatalogTableTableExtractor().apply(spark, ct)
     }

diff --git a/...ubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala b/...ubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/AuthZUtils.scala
@@ -62,10 +62,23 @@ private[authz] object AuthZUtils {
       obj: AnyRef,
       methodName: String,
       args: (Class[_], AnyRef)*): AnyRef = {
-    val (types, values) = args.unzip
-    val method = obj.getClass.getMethod(methodName, types: _*)
-    method.setAccessible(true)
-    method.invoke(obj, values: _*)
+    try {
+      val (types, values) = args.unzip
+      val method = obj.getClass.getMethod(methodName, types: _*)
+      method.setAccessible(true)
+      method.invoke(obj, values: _*)
+    } catch {
+      case e: NoSuchMethodException =>
+        val candidates = obj.getClass.getMethods.map(_.getName).mkString("[", ",", "]")
+        throw new RuntimeException(s"$methodName not in ${obj.getClass} $candidates", e)
+    }
+  }
+
+  def invokeAs[T](
+      obj: AnyRef,
+      methodName: String,
+      args: (Class[_], AnyRef)*): T = {
+    invoke(obj, methodName, args: _*).asInstanceOf[T]
   }
 
   def invokeStatic(

diff --git a/...authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/RangerConfigProvider.scala b/...authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/util/RangerConfigProvider.scala
@@ -36,7 +36,7 @@ trait RangerConfigProvider {
   def getRangerConf: Configuration = {
     try {
       // for Ranger 2.1+
-      invoke(this, "getConfig").asInstanceOf[Configuration]
+      invokeAs[Configuration](this, "getConfig")
     } catch {
       case _: NoSuchMethodException =>
         // for Ranger 2.0 and below