diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala index cb233b4651f..38211408415 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarker.scala @@ -17,11 +17,13 @@ package org.apache.kyuubi.plugin.spark.authz.rule.permanentview +import scala.annotation.tailrec + import org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation import org.apache.spark.sql.catalyst.catalog.CatalogTable -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Cast} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute} import org.apache.spark.sql.catalyst.plans.QueryPlan -import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan, Project, Statistics, View} +import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan, Project, Statistics} import org.apache.spark.sql.catalyst.trees.TreeNodeTag case class PermanentViewMarker(child: LogicalPlan, catalogTable: CatalogTable) @@ -39,7 +41,7 @@ case class PermanentViewMarker(child: LogicalPlan, catalogTable: CatalogTable) override def newInstance(): LogicalPlan = { val projectList = child.output.map { case attr => - Alias(Cast(attr, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata)) + Alias(attr, attr.name)(explicitMetadata = Some(attr.metadata)) } val newProj = Project(projectList, child) newProj.setTagValue(PVM_NEW_INSTANCE_TAG, ()) @@ -48,11 +50,14 @@ case class PermanentViewMarker(child: LogicalPlan, catalogTable: CatalogTable) } override def doCanonicalize(): LogicalPlan = { - child match { - case p @ Project(_, view: View) if p.getTagValue(PVM_NEW_INSTANCE_TAG).contains(true) => - view.canonicalized - case _ => - child.canonicalized + // newInstance() wraps the child in a Project, and a new instance of a new instance nests one + // inside another, so strip every layer it added rather than only the outermost one. + @tailrec + def stripNewInstanceProjects(plan: LogicalPlan): LogicalPlan = plan match { + case p: Project if p.getTagValue(PVM_NEW_INSTANCE_TAG).isDefined => + stripNewInstanceProjects(p.child) + case other => other } + stripNewInstanceProjects(child).canonicalized } } 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 8805f0cd527..401ae075e4f 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 @@ -741,6 +741,42 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { } } + test("[KYUUBI #5937] cached permanent view is reused when referenced twice in one query") { + val db1 = defaultDb + val srcTable = "pvm_cache_src" + val view = "pvm_cache_view" + + withCleanTmpResources(Seq((s"$db1.$view", "view"), (s"$db1.$srcTable", "table"))) { + doAs(admin, sql(s"CREATE TABLE IF NOT EXISTS $db1.$srcTable (id int, k int)")) + doAs(admin, sql(s"CREATE OR REPLACE VIEW $db1.$view AS SELECT id, k FROM $db1.$srcTable")) + + doAs(admin) { + val df = spark.table(s"$db1.$view") + df.cache() + try { + df.count() + // The same cached relation appears twice in one plan, so the analyzer replaces one of + // them with PermanentViewMarker.newInstance(). Both occurrences should still be served + // from the cache. + val optimized = df.join(df.select("k").distinct(), "k").queryExecution.optimizedPlan + assert(optimized.collect { case r: InMemoryRelation => r }.size === 2) + assert(optimized.collect { case r: HiveTableRelation => r }.isEmpty) + assert(optimized.collect { case r: LogicalRelation => r }.isEmpty) + + // Reusing an already analysed plan feeds a renewed marker back in, so newInstance() + // nests one Project inside another. Every layer has to be seen through. + val reused = df.join(df.select("k").distinct(), "k") + val nested = reused.join(reused.select("k").distinct(), "k") + .queryExecution.optimizedPlan + assert(nested.collect { case r: HiveTableRelation => r }.isEmpty) + assert(nested.collect { case r: LogicalRelation => r }.isEmpty) + } finally { + df.unpersist() + } + } + } + } + test("[KYUUBI #3608] Support {OWNER} variable for queries") { val db = defaultDb val table = "owner_variable" diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarkerSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarkerSuite.scala new file mode 100644 index 00000000000..216ec8d5dd1 --- /dev/null +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/rule/permanentview/PermanentViewMarkerSuite.scala @@ -0,0 +1,73 @@ +/* + * 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.permanentview + +import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType} +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Project, View} +import org.apache.spark.sql.types.StructType + +import org.apache.kyuubi.KyuubiFunSuite + +class PermanentViewMarkerSuite extends KyuubiFunSuite { + + private def newMarker(): PermanentViewMarker = { + val desc = CatalogTable( + identifier = TableIdentifier("v", Some("default")), + tableType = CatalogTableType.VIEW, + storage = CatalogStorageFormat.empty, + schema = new StructType().add("a", "int").add("b", "string"), + viewText = Some("SELECT a, b FROM t")) + val child: LogicalPlan = View(desc, isTempView = false, LocalRelation($"a".int, $"b".string)) + PermanentViewMarker(child, desc) + } + + test("new instance canonicalizes to the same plan as the original") { + // PermanentViewMarker is a MultiInstanceRelation, so the analyzer may replace an occurrence + // with newInstance() whenever the same relation appears twice in one plan. A new instance has + // to keep sameResult with the original, otherwise CacheManager stops recognising it and the + // view is read from its sources again. + val marker = newMarker() + assert(marker.newInstance().canonicalized == marker.canonicalized) + } + + test("nested new instances canonicalize to the same plan as the original") { + // newInstance() wraps the child in a Project, so calling it on a plan that is already a new + // instance nests one Project inside another. Canonicalization has to see through every layer, + // not just the first one. + val marker = newMarker() + val once = marker.newInstance().asInstanceOf[PermanentViewMarker] + assert(once.newInstance().canonicalized == marker.canonicalized) + } + + test("a marker over a plan that is not a View canonicalizes to the same plan") { + // RuleApplyPermanentViewMarker also wraps the plan of every SubqueryExpression found inside + // a view, and that plan is not a View. Canonicalization has to cover those markers too. + val desc = CatalogTable( + identifier = TableIdentifier("v", Some("default")), + tableType = CatalogTableType.VIEW, + storage = CatalogStorageFormat.empty, + schema = new StructType().add("a", "int").add("b", "string"), + viewText = Some("SELECT a, b FROM t WHERE a IN (SELECT a FROM s)")) + val subqueryPlan: LogicalPlan = + Project(Seq($"a".int, $"b".string), LocalRelation($"a".int, $"b".string)) + val marker = PermanentViewMarker(subqueryPlan, desc) + assert(marker.newInstance().canonicalized == marker.canonicalized) + } +}