[incubator-kyuubi] branch master updated: [KYUUBI #3226] Privileges should be checked only once in `RuleAuthorization`

[ul...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ulyssesyou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 98df59107 [KYUUBI #3226] Privileges should be checked only once in `RuleAuthorization`
98df59107 is described below

commit 98df59107d00c60d612444606fb15229e88086d0
Author: yikf <yi...@gmail.com>
AuthorDate: Thu Aug 18 17:58:03 2022 +0800

    [KYUUBI #3226] Privileges should be checked only once in `RuleAuthorization`
    
    ### _Why are the changes needed?_
    
    fix https://github.com/apache/incubator-kyuubi/issues/3226
    
    ### _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
    
    - [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #3262 from yikf/auth-once.
    
    Closes #3226
    
    59c1321e [yikf] fix
    dad11aa8 [yikf] [KYUUBI #3226] check privileges once only
    
    Authored-by: yikf <yi...@gmail.com>
    Signed-off-by: ulysses-you <ul...@apache.org>
---
 .../spark/authz/ranger/RuleAuthorization.scala     | 13 ++++-
 .../authz/ranger/RangerSparkExtensionSuite.scala   | 64 ++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala
index c22fca632..4a0b95ccf 100644
--- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala
+++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RuleAuthorization.scala
@@ -22,20 +22,27 @@ import scala.collection.mutable.ArrayBuffer
 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.{ObjectType, _}
 import org.apache.kyuubi.plugin.spark.authz.ObjectType._
+import org.apache.kyuubi.plugin.spark.authz.ranger.RuleAuthorization.KYUUBI_AUTHZ_TAG
 import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._
 
 class RuleAuthorization(spark: SparkSession) extends Rule[LogicalPlan] {
-  override def apply(plan: LogicalPlan): LogicalPlan = {
-    RuleAuthorization.checkPrivileges(spark, plan)
-    plan
+  override def apply(plan: LogicalPlan): LogicalPlan = plan match {
+    case p if !plan.getTagValue(KYUUBI_AUTHZ_TAG).contains(true) =>
+      RuleAuthorization.checkPrivileges(spark, p)
+      p.setTagValue(KYUUBI_AUTHZ_TAG, true)
+      p
+    case p => p // do nothing if checked privileges already.
   }
 }
 
 object RuleAuthorization {
 
+  val KYUUBI_AUTHZ_TAG = TreeNodeTag[Boolean]("__KYUUBI_AUTHZ_TAG")
+
   def checkPrivileges(spark: SparkSession, plan: LogicalPlan): Unit = {
     val auditHandler = new SparkRangerAuditHandler
     val ugi = getAuthzUgi(spark.sparkContext)
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 ec48c6587..ae519b64c 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
@@ -27,12 +27,14 @@ import org.apache.hadoop.security.UserGroupInformation
 import org.apache.spark.sql.{Row, SparkSessionExtensions}
 import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
 import org.apache.spark.sql.catalyst.plans.logical.Statistics
+import org.apache.spark.sql.execution.columnar.InMemoryRelation
 import org.apache.spark.sql.execution.datasources.LogicalRelation
 import org.scalatest.BeforeAndAfterAll
 // scalastyle:off
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.kyuubi.plugin.spark.authz.SparkSessionProvider
+import org.apache.kyuubi.plugin.spark.authz.ranger.RuleAuthorization.KYUUBI_AUTHZ_TAG
 import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils.getFieldVal
 
 abstract class RangerSparkExtensionSuite extends AnyFunSuite
@@ -59,6 +61,68 @@ abstract class RangerSparkExtensionSuite extends AnyFunSuite
     s"Permission denied: user [$user] does not have [$privilege] privilege on [$resource]"
   }
 
+  test("[KYUUBI #3226] RuleAuthorization: Should check privileges once only.") {
+    val logicalPlan = doAs("admin", sql("SHOW TABLES").queryExecution.logical)
+    val rule = new RuleAuthorization(spark)
+
+    (1 until 10).foreach { i =>
+      if (i == 1) {
+        assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).isEmpty)
+      } else {
+        assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+      }
+      rule.apply(logicalPlan)
+    }
+
+    assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+  }
+
+  test("[KYUUBI #3226]: Another session should also check even if the plan is cached.") {
+    val testTable = "mytable"
+    val create =
+      s"""
+         | CREATE TABLE IF NOT EXISTS $testTable
+         | (id string)
+         | using parquet
+         |""".stripMargin
+    val select = s"SELECT * FROM $testTable"
+    try {
+      // create tmp table
+      doAs(
+        "admin", {
+          sql(create)
+
+          // session1: first query, should auth once.[LogicalRelation]
+          val df = sql(select)
+          val plan1 = df.queryExecution.optimizedPlan
+          assert(plan1.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+
+          // cache
+          df.cache()
+
+          // session1: second query, should auth once.[InMemoryRelation]
+          // (don't need to check in again, but it's okay to check in once)
+          val plan2 = sql(select).queryExecution.optimizedPlan
+          assert(plan1 != plan2 && plan2.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+
+          // session2: should auth once.
+          val otherSessionDf = spark.newSession().sql(select)
+
+          // KYUUBI_AUTH_TAG is None
+          assert(otherSessionDf.queryExecution.logical.getTagValue(KYUUBI_AUTHZ_TAG).isEmpty)
+          val plan3 = otherSessionDf.queryExecution.optimizedPlan
+          // make sure it use cache.
+          assert(plan3.isInstanceOf[InMemoryRelation])
+          // auth once only.
+          assert(plan3.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+        })
+
+    } finally {
+      // finally, drop tmp table
+      doAs("admin", sql(s"DROP TABLE IF EXISTS $testTable"))
+    }
+  }
+
   test("auth: databases") {
     val testDb = "mydb"
     val create = s"CREATE DATABASE IF NOT EXISTS $testDb"