From 62fb9cdfc29fea80235330d65776c556dc34e89e Mon Sep 17 00:00:00 2001
From: Angerszhuuuu <angers.zhu@gmail.com>
Date: Mon, 25 Dec 2023 14:34:52 +0800
Subject: [PATCH] [KYUUBI #5913][Bug] After resolve PVM should mark all nodes
 as checked
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# :mag: Description
## Issue References 🔗

This pull request fixes #5913

## Describe Your Solution 🔧

We meet a case that seem after `RuleEliminatePermanentViewMarker` apply and optimize the PVM's child plan, seems some case node's tag was missed, then also check the PVM's source table again when `OptimizeSubqueries`.

We should mark all node as checked for pvm's child and optimized child.
1. It's more stable
2. It's safe

## Types of changes :bookmark:

- [ ] Bugfix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)

## Test Plan 🧪

Manuel tested in our prod, but didn't reproduce it in the UT since the case is too complex SQL.
#### Behavior Without This Pull Request :coffin:

#### Behavior With This Pull Request :tada:

#### Related Unit Tests

---

# Checklists
## 📝 Author Self Checklist

- [x] My code follows the [style guidelines](https://kyuubi.readthedocs.io/en/master/contributing/code/style.html) of this project
- [x] I have performed a self-review
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html)

## 📝 Committer Pre-Merge Checklist

- [x] Pull request title is okay.
- [x] No license issues.
- [x] Milestone correctly set?
- [x] Test coverage is ok
- [x] Assignees are selected.
- [x] Minimum number of approvals
- [x] No changes are requested

**Be nice. Be informative.**

Closes #5915 from AngersZhuuuu/KYUUBI-5913.

Closes #5913

38c04bd70 [Angerszhuuuu] [KYUUBI #5913][Bug] After resolve PVM should mark all nodes as checke

Authored-by: Angerszhuuuu <angers.zhu@gmail.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../kyuubi/plugin/spark/authz/rule/Authorization.scala     | 2 +-
 .../authz/rule/RuleEliminatePermanentViewMarker.scala      | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala
index d1494266e..d682b71d9 100644
--- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala
+++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/Authorization.scala
@@ -43,7 +43,7 @@ object Authorization {
 
   val KYUUBI_AUTHZ_TAG = TreeNodeTag[Unit]("__KYUUBI_AUTHZ_TAG")
 
-  private def markAllNodesAuthChecked(plan: LogicalPlan): LogicalPlan = {
+  def markAllNodesAuthChecked(plan: LogicalPlan): LogicalPlan = {
     plan.transformDown { case p =>
       p.setTagValue(KYUUBI_AUTHZ_TAG, ())
       p
diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala
index 003521c72..a0a6d5b62 100644
--- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala
+++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminatePermanentViewMarker.scala
@@ -37,7 +37,7 @@ case class RuleEliminatePermanentViewMarker(sparkSession: SparkSession) extends
         }
         // For each SubqueryExpression's PVM, we should mark as resolved to
         // avoid check privilege of PVM's internal Subquery.
-        Authorization.markAuthChecked(ret)
+        Authorization.markAllNodesAuthChecked(ret)
         ret
     }
   }
@@ -52,8 +52,9 @@ case class RuleEliminatePermanentViewMarker(sparkSession: SparkSession) extends
         }
     }
     if (matched) {
-      Authorization.markAuthChecked(eliminatedPVM)
-      sparkSession.sessionState.optimizer.execute(eliminatedPVM)
+      Authorization.markAllNodesAuthChecked(eliminatedPVM)
+      val optimized = sparkSession.sessionState.optimizer.execute(eliminatedPVM)
+      Authorization.markAllNodesAuthChecked(optimized)
     } else {
       eliminatedPVM
     }