update

zhaohehuhu · Nov 27, 2023 · 8e53236 · 8e53236
1 parent 3bafa7c
commit 8e53236
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 42 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
@@ -28,7 +28,6 @@ import org.slf4j.LoggerFactory
 import org.apache.kyuubi.plugin.spark.authz.OperationType.OperationType
 import org.apache.kyuubi.plugin.spark.authz.PrivilegeObjectActionType._
 import org.apache.kyuubi.plugin.spark.authz.rule.Authorization._
-import org.apache.kyuubi.plugin.spark.authz.rule.permanentview.PermanentViewMarker
 import org.apache.kyuubi.plugin.spark.authz.serde._
 import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._
 import org.apache.kyuubi.util.reflect.ReflectUtils._
@@ -67,12 +66,7 @@ object PrivilegesBuilder {
 
     def mergeProjection(table: Table, plan: LogicalPlan): Unit = {
       if (projectionList.isEmpty) {
-        plan match {
-          case pvm: PermanentViewMarker =>
-            privilegeObjects += PrivilegeObject(table, pvm.outputColNames)
-          case _ =>
-            privilegeObjects += PrivilegeObject(table, plan.output.map(_.name))
-        }
+        privilegeObjects += PrivilegeObject(table, plan.output.map(_.name))
       } else {
         val cols = columnPrune(projectionList ++ conditionList, plan.outputSet)
         privilegeObjects += PrivilegeObject(table, cols)
@@ -137,22 +131,16 @@ object PrivilegesBuilder {
 
       case p =>
         for (child <- p.children) {
-          child match {
-            case pvm: PermanentViewMarker
-                if pvm.isSubqueryExpressionPlaceHolder || pvm.output.isEmpty =>
-              buildQuery(child, privilegeObjects, projectionList, conditionList, spark)
-            case _ =>
-              val childCols = columnPrune(projectionList ++ conditionList, child.outputSet)
-              if (childCols.isEmpty) {
-                buildQuery(
-                  child,
-                  privilegeObjects,
-                  p.inputSet.map(_.toAttribute).toSeq,
-                  conditionList,
-                  spark)
-              } else {
-                buildQuery(child, privilegeObjects, projectionList, conditionList, spark)
-              }
+          val childCols = columnPrune(projectionList ++ conditionList, child.outputSet)
+          if (childCols.isEmpty) {
+            buildQuery(
+              child,
+              privilegeObjects,
+              p.inputSet.map(_.toAttribute).toSeq,
+              conditionList,
+              spark)
+          } else {
+            buildQuery(child, privilegeObjects, projectionList, conditionList, spark)
           }
         }
     }

diff --git a/...thz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala b/...thz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala
@@ -53,7 +53,7 @@ class RangerSparkExtension extends (SparkSessionExtensions => Unit) {
     v1.injectResolutionRule(RuleApplyDataMaskingStage1)
     v1.injectOptimizerRule(_ => new RuleEliminateMarker())
     v1.injectOptimizerRule(new RuleAuthorization(_))
-    v1.injectOptimizerRule(_ => new RuleEliminatePermanentViewMarker())
+    v1.injectOptimizerRule(new RuleEliminatePermanentViewMarker(_))
     v1.injectOptimizerRule(_ => new RuleEliminateTypeOf())
     v1.injectPlannerStrategy(new FilterDataSourceV2Strategy(_))
   }

diff --git a/...-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala b/...-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala
@@ -51,7 +51,7 @@ object Authorization {
     }
   }
 
-  protected def markAuthChecked(plan: LogicalPlan): LogicalPlan = {
+  def markAuthChecked(plan: LogicalPlan): LogicalPlan = {
     plan.setTagValue(KYUUBI_AUTHZ_TAG, ())
     plan transformDown {
       case pvm: PermanentViewMarker =>

diff --git a/...in/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala b/...in/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala
@@ -17,6 +17,7 @@
 
 package org.apache.kyuubi.plugin.spark.authz.rule
 
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -26,12 +27,21 @@ import org.apache.kyuubi.plugin.spark.authz.rule.permanentview.PermanentViewMark
 /**
  * Transforming up [[PermanentViewMarker]]
  */
-class RuleEliminatePermanentViewMarker extends Rule[LogicalPlan] {
+class RuleEliminatePermanentViewMarker(sparkSession: SparkSession) extends Rule[LogicalPlan] {
   override def apply(plan: LogicalPlan): LogicalPlan = {
-    plan.transformUp {
-      case pvm: PermanentViewMarker => pvm.child.transformAllExpressions {
+    var matched = false
+    val eliminatedPVM = plan.transformUp {
+      case pvm: PermanentViewMarker =>
+        matched = true
+        pvm.child.transformAllExpressions {
           case s: SubqueryExpression => s.withNewPlan(apply(s.plan))
         }
     }
+    if (matched) {
+      Authorization.markAuthChecked(eliminatedPVM)
+      sparkSession.sessionState.optimizer.execute(eliminatedPVM)
+    } else {
+      eliminatedPVM
+    }
   }
 }
diff --git a/...n/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala b/...n/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala
@@ -23,11 +23,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
 
 import org.apache.kyuubi.plugin.spark.authz.util.WithInternalChild
 
-case class PermanentViewMarker(
-    child: LogicalPlan,
-    catalogTable: CatalogTable,
-    outputColNames: Seq[String],
-    isSubqueryExpressionPlaceHolder: Boolean = false) extends LeafNode
+case class PermanentViewMarker(child: LogicalPlan, catalogTable: CatalogTable) extends LeafNode
   with WithInternalChild {
 
   override def output: Seq[Attribute] = child.output

diff --git a/...rg/apache/kyuubi/plugin/spark/authz/rule/permanentview/RuleApplyPermanentViewMarker.scala b/...rg/apache/kyuubi/plugin/spark/authz/rule/permanentview/RuleApplyPermanentViewMarker.scala
@@ -36,16 +36,11 @@ class RuleApplyPermanentViewMarker extends Rule[LogicalPlan] {
     plan mapChildren {
       case p: PermanentViewMarker => p
       case permanentView: View if hasResolvedPermanentView(permanentView) =>
-        val resolved = permanentView.transformAllExpressions {
+        val resolved = permanentView.child.transformAllExpressions {
           case subquery: SubqueryExpression =>
-            subquery.withNewPlan(plan =
-              PermanentViewMarker(
-                subquery.plan,
-                permanentView.desc,
-                permanentView.output.map(_.name),
-                true))
+            subquery.withNewPlan(plan = PermanentViewMarker(subquery.plan, permanentView.desc))
         }
-        PermanentViewMarker(resolved, resolved.desc, resolved.output.map(_.name))
+        PermanentViewMarker(resolved, permanentView.desc)
       case other => apply(other)
     }
   }

diff --git a/...rc/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala b/...rc/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala
@@ -892,7 +892,7 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite {
             sql(s"SELECT id as new_id, name, max_scope FROM $db1.$view1".stripMargin).show()))
         assert(e2.getMessage.contains(
           s"does not have [select] privilege on " +
-            s"[$db1/$view1/id,$db1/$view1/name,$db1/$view1/max_scope,$db1/$view1/sum_age]"))
+            s"[$db1/$view1/id,$db1/$view1/name,$db1/$view1/max_scope]"))
       }
     }
   }