From 8e53236ac6bb675f25eec2245412202b3ba06b84 Mon Sep 17 00:00:00 2001 From: Angerszhuuuu Date: Mon, 27 Nov 2023 14:45:31 +0800 Subject: [PATCH] update --- .../spark/authz/PrivilegesBuilder.scala | 34 ++++++------------- .../authz/ranger/RangerSparkExtension.scala | 2 +- .../spark/authz/rule/Authorization.scala | 2 +- .../RuleEliminatePermanentViewMarker.scala | 16 +++++++-- .../permanentview/PermanentViewMarker.scala | 6 +--- .../RuleApplyPermanentViewMarker.scala | 11 ++---- .../ranger/RangerSparkExtensionSuite.scala | 2 +- 7 files changed, 31 insertions(+), 42 deletions(-) diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala index 8bd29f110ac..c0fe8db370c 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala +++ b/extensions/spark/kyuubi-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/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala index f6d43900370..93c10068a8c 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala +++ b/extensions/spark/kyuubi-spark-authz/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/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala index d83541825e4..cbe05478a40 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala +++ b/extensions/spark/kyuubi-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/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala index 864ada55f94..d468dcca614 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/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/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala index d7f7e344147..5ef8f4a748d 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/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/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/RuleApplyPermanentViewMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/RuleApplyPermanentViewMarker.scala index 1a0024abb7e..87d7948af2f 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/RuleApplyPermanentViewMarker.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/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/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala index 66a5ef05849..a57d86e63f5 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/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]")) } } }