[KYUUBI #5555][AUTHZ] Extractor common authorization rule class

apache · Oct 30, 2023 · 623b351 · 623b351
1 parent 26f614a
commit 623b351
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 36 deletions.
diff --git a/...-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala b/...-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala
@@ -22,30 +22,15 @@ import scala.collection.mutable.ArrayBuffer
 import org.apache.ranger.plugin.policyengine.RangerAccessRequest
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 
 import org.apache.kyuubi.plugin.spark.authz._
 import org.apache.kyuubi.plugin.spark.authz.ObjectType._
-import org.apache.kyuubi.plugin.spark.authz.ranger.RuleAuthorization._
 import org.apache.kyuubi.plugin.spark.authz.ranger.SparkRangerAdminPlugin._
-import org.apache.kyuubi.plugin.spark.authz.rule.permanentview.PermanentViewMarker
+import org.apache.kyuubi.plugin.spark.authz.rule.Authorization
 import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._
 
-class RuleAuthorization(spark: SparkSession) extends Rule[LogicalPlan] {
-  override def apply(plan: LogicalPlan): LogicalPlan = {
-    plan match {
-      case plan if isAuthChecked(plan) => plan // do nothing if checked privileges already.
-      case p => checkPrivileges(spark, p)
-    }
-  }
-}
-
-object RuleAuthorization {
-
-  val KYUUBI_AUTHZ_TAG = TreeNodeTag[Unit]("__KYUUBI_AUTHZ_TAG")
-
-  private def checkPrivileges(spark: SparkSession, plan: LogicalPlan): LogicalPlan = {
+class RuleAuthorization(spark: SparkSession) extends Authorization(spark) {
+  override def checkPrivileges(spark: SparkSession, plan: LogicalPlan): Unit = {
     val auditHandler = new SparkRangerAuditHandler
     val ugi = getAuthzUgi(spark.sparkContext)
     val (inputs, outputs, opType) = PrivilegesBuilder.build(plan, spark)
@@ -95,23 +80,5 @@ object RuleAuthorization {
         verify(Seq(req), auditHandler)
       }
     }
-    markAuthChecked(plan)
-  }
-
-  private def markAuthChecked(plan: LogicalPlan): LogicalPlan = {
-    plan match {
-      case _: PermanentViewMarker =>
-        plan.transformUp { case p =>
-          p.setTagValue(KYUUBI_AUTHZ_TAG, ())
-          p
-        }
-      case _ =>
-        plan.setTagValue(KYUUBI_AUTHZ_TAG, ())
-    }
-    plan
-  }
-
-  private def isAuthChecked(plan: LogicalPlan): Boolean = {
-    plan.find(_.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty).nonEmpty
   }
 }
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
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.plugin.spark.authz.rule
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+
+import org.apache.kyuubi.plugin.spark.authz.rule.Authorization._
+import org.apache.kyuubi.plugin.spark.authz.rule.permanentview.PermanentViewMarker
+
+abstract class Authorization(spark: SparkSession) extends Rule[LogicalPlan] {
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    plan match {
+      case plan if isAuthChecked(plan) => plan // do nothing if checked privileges already.
+      case p =>
+        checkPrivileges(spark, p)
+        markAuthChecked(p)
+    }
+  }
+
+  def checkPrivileges(spark: SparkSession, plan: LogicalPlan): Unit
+}
+
+object Authorization {
+
+  val KYUUBI_AUTHZ_TAG = TreeNodeTag[Unit]("__KYUUBI_AUTHZ_TAG")
+
+  protected def markAuthChecked(plan: LogicalPlan): LogicalPlan = {
+    plan match {
+      case _: PermanentViewMarker =>
+        plan.transformUp { case p =>
+          p.setTagValue(KYUUBI_AUTHZ_TAG, ())
+          p
+        }
+      case _ =>
+        plan.setTagValue(KYUUBI_AUTHZ_TAG, ())
+    }
+    plan
+  }
+
+  protected def isAuthChecked(plan: LogicalPlan): Boolean = {
+    plan.find(_.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty).nonEmpty
+  }
+}