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 d1494266e..d682b71d9 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 @@ -43,7 +43,7 @@ object Authorization { val KYUUBI_AUTHZ_TAG = TreeNodeTag[Unit]("__KYUUBI_AUTHZ_TAG") - private def markAllNodesAuthChecked(plan: LogicalPlan): LogicalPlan = { + def markAllNodesAuthChecked(plan: LogicalPlan): LogicalPlan = { plan.transformDown { case p => p.setTagValue(KYUUBI_AUTHZ_TAG, ()) p 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 003521c72..a0a6d5b62 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 @@ -37,7 +37,7 @@ case class RuleEliminatePermanentViewMarker(sparkSession: SparkSession) extends } // For each SubqueryExpression's PVM, we should mark as resolved to // avoid check privilege of PVM's internal Subquery. - Authorization.markAuthChecked(ret) + Authorization.markAllNodesAuthChecked(ret) ret } } @@ -52,8 +52,9 @@ case class RuleEliminatePermanentViewMarker(sparkSession: SparkSession) extends } } if (matched) { - Authorization.markAuthChecked(eliminatedPVM) - sparkSession.sessionState.optimizer.execute(eliminatedPVM) + Authorization.markAllNodesAuthChecked(eliminatedPVM) + val optimized = sparkSession.sessionState.optimizer.execute(eliminatedPVM) + Authorization.markAllNodesAuthChecked(optimized) } else { eliminatedPVM }