[KYUUBI #5594][AUTHZ] BuildQuery should respect normal node's input

extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala

@@ -20,7 +20,7 @@ package org.apache.kyuubi.plugin.spark.authz
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression}
+import org.apache.spark.sql.catalyst.expressions.{AttributeSet, Expression, NamedExpression}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.command.ExplainCommand
 import org.slf4j.LoggerFactory
@@ -69,12 +69,19 @@ object PrivilegesBuilder {
       if (projectionList.isEmpty) {
         privilegeObjects += PrivilegeObject(table, plan.output.map(_.name))
       } else {
-        val cols = (projectionList ++ conditionList).flatMap(collectLeaves)
-          .filter(plan.outputSet.contains).map(_.name).distinct
+        val cols = columnPrune(projectionList ++ conditionList, plan.outputSet)
         privilegeObjects += PrivilegeObject(table, cols)
       }
     }
 
+    def columnPrune(projectionList: Seq[NamedExpression], output: AttributeSet): Seq[String] = {
+      (projectionList ++ conditionList)
+        .flatMap(collectLeaves)
+        .filter(output.contains)
+        .map(_.name)
+        .distinct
+    }
+
     plan match {
       case p if p.getTagValue(KYUUBI_AUTHZ_TAG).nonEmpty =>
 
@@ -125,7 +132,17 @@ object PrivilegesBuilder {
 
       case p =>
         for (child <- p.children) {
-          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)
+          }
         }
     }
   }

extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala

@@ -23,12 +23,15 @@ import java.nio.file.Path
 import scala.util.Try
 
 import org.apache.hadoop.security.UserGroupInformation
-import org.apache.spark.sql.{Row, SparkSessionExtensions}
+import org.apache.spark.sql.{DataFrame, Row, SparkSessionExtensions}
 import org.apache.spark.sql.catalyst.analysis.NoSuchTableException
 import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
 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.apache.spark.sql.functions.col
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 import org.scalatest.BeforeAndAfterAll
 // scalastyle:off
 import org.scalatest.funsuite.AnyFunSuite
@@ -1368,4 +1371,80 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite {
       }
     }
   }
+
+  test("[KYUUBI #5594][AUTHZ] BuildQuery should respect normal node's input ") {
+    assume(!isSparkV35OrGreater, "mapInPandas not supported after spark 3.5")
+    val db1 = defaultDb
+    val table1 = "table1"
+    val view1 = "view1"
+    withSingleCallEnabled {
+      withCleanTmpResources(Seq((s"$db1.$table1", "table"), (s"$db1.$view1", "view"))) {
+        doAs(admin, sql(s"CREATE TABLE IF NOT EXISTS $db1.$table1 (id int, scope int)"))
+        doAs(admin, sql(s"CREATE VIEW $db1.$view1 AS SELECT * FROM $db1.$table1"))
+
+        val table = spark.read.table(s"$db1.$table1")
+        val mapTableInPandasUDF = PythonUDF(
+          "mapInPandasUDF",
+          null,
+          StructType(Seq(StructField("id", IntegerType), StructField("scope", IntegerType))),
+          table.queryExecution.analyzed.output,
+          205,
+          true)
+        interceptContains[AccessControlException](
+          doAs(
+            someone,
+            invokeAs(
+              table,
+              "mapInPandas",
+              (classOf[PythonUDF], mapTableInPandasUDF))
+              .asInstanceOf[DataFrame].select(col("id"), col("scope")).limit(1).show(true)))(
+          s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]")
+
+        val view = spark.read.table(s"$db1.$view1")
+        val mapViewInPandasUDF = PythonUDF(
+          "mapInPandasUDF",
+          null,
+          StructType(Seq(StructField("id", IntegerType), StructField("scope", IntegerType))),
+          view.queryExecution.analyzed.output,
+          205,
+          true)
+        interceptContains[AccessControlException](
+          doAs(
+            someone,
+            invokeAs(
+              view,
+              "mapInPandas",
+              (classOf[PythonUDF], mapViewInPandasUDF))
+              .asInstanceOf[DataFrame].select(col("id"), col("scope")).limit(1).show(true)))(
+          s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]")
+      }
+    }
+  }
+
+  test("[KYUUBI #5594][AUTHZ] BuildQuery should respect sort agg input") {
+    val db1 = defaultDb
+    val table1 = "table1"
+    val view1 = "view1"
+    withSingleCallEnabled {
+      withCleanTmpResources(Seq((s"$db1.$table1", "table"), (s"$db1.$view1", "view"))) {
+        doAs(admin, sql(s"CREATE TABLE IF NOT EXISTS $db1.$table1 (id int, scope int)"))
+        doAs(admin, sql(s"CREATE VIEW $db1.$view1 AS SELECT * FROM $db1.$table1"))
+        interceptContains[AccessControlException](
+          doAs(
+            someone,
+            sql(s"SELECT count(*) FROM $db1.$table1 WHERE id > 1").collect()))(
+          s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]")
+        interceptContains[AccessControlException](
+          doAs(
+            someone,
+            sql(s"SELECT count(*) FROM $db1.$view1 WHERE id > 1").collect()))(
+          s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]")
+        interceptContains[AccessControlException](
+          doAs(
+            someone,
+            sql(s"SELECT count(id) FROM $db1.$view1 WHERE id > 1").collect()))(
+          s"does not have [select] privilege on [$db1/$view1/id]")
+      }
+    }
+  }
 }