follow comment

apache · Oct 27, 2023 · f6b6dcc · f6b6dcc
1 parent 2c7ae5d
commit f6b6dcc
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 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
@@ -42,7 +42,7 @@ class RuleAuthorization(spark: SparkSession) extends Rule[LogicalPlan] {
 
 object RuleAuthorization {
 
-  val KYUUBI_AUTHZ_TAG = TreeNodeTag[Boolean]("__KYUUBI_AUTHZ_TAG")
+  val KYUUBI_AUTHZ_TAG = TreeNodeTag[Unit]("__KYUUBI_AUTHZ_TAG")
 
   private def checkPrivileges(spark: SparkSession, plan: LogicalPlan): LogicalPlan = {
     val auditHandler = new SparkRangerAuditHandler
@@ -101,16 +101,16 @@ object RuleAuthorization {
     plan match {
       case _: PermanentViewMarker =>
         plan.transformUp { case p =>
-          p.setTagValue(KYUUBI_AUTHZ_TAG, true)
+          p.setTagValue(KYUUBI_AUTHZ_TAG, ())
           p
         }
       case _ =>
-        plan.setTagValue(KYUUBI_AUTHZ_TAG, true)
+        plan.setTagValue(KYUUBI_AUTHZ_TAG, ())
     }
     plan
   }
 
   private def isAuthChecked(plan: LogicalPlan): Boolean = {
-    plan.find(_.getTagValue(KYUUBI_AUTHZ_TAG).contains(true)).nonEmpty
+    plan.find(_.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty).nonEmpty
   }
 }
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
@@ -113,12 +113,12 @@ abstract class RangerSparkExtensionSuite extends AnyFunSuite
       if (i == 1) {
         assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).isEmpty)
       } else {
-        assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+        assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty)
       }
       rule.apply(logicalPlan)
     }
 
-    assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+    assert(logicalPlan.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty)
   }
 
   test("[KYUUBI #3226]: Another session should also check even if the plan is cached.") {
@@ -140,15 +140,15 @@ abstract class RangerSparkExtensionSuite extends AnyFunSuite
           // session1: first query, should auth once.[LogicalRelation]
           val df = sql(select)
           val plan1 = df.queryExecution.optimizedPlan
-          assert(plan1.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+          assert(plan1.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty)
 
           // 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))
+          assert(plan1 != plan2 && plan2.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty)
 
           // session2: should auth once.
           val otherSessionDf = spark.newSession().sql(select)
@@ -159,7 +159,7 @@ abstract class RangerSparkExtensionSuite extends AnyFunSuite
           // make sure it use cache.
           assert(plan3.isInstanceOf[InMemoryRelation])
           // auth once only.
-          assert(plan3.getTagValue(KYUUBI_AUTHZ_TAG).getOrElse(false))
+          assert(plan3.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty)
         })
     }
   }