From 52b238edc235d3f207be6840da513a79682d8f11 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Mon, 6 Jan 2025 15:45:09 +0800
Subject: [PATCH v1 2/2] refactoring LET statement self assign privilege check

---
 src/backend/optimizer/plan/planner.c |  9 ++++---
 src/backend/optimizer/plan/setrefs.c | 37 ++++++++++++++++++++--------
 src/include/nodes/parsenodes.h       |  9 +++++++
 src/include/nodes/pathnodes.h        |  6 -----
 4 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 11f736b437..a21d8e48c6 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -334,8 +334,6 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	glob->transientPlan = false;
 	glob->dependsOnRole = false;
 	glob->sessionVariables = NIL;
-	glob->basenodeSessionVarid = InvalidOid;
-	glob->checkSelectPermVarids = NULL;
 
 	/*
 	 * Assess whether it's feasible to use parallel mode for this query. We
@@ -578,8 +576,11 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	 * for assignemnt indirection should be excluded from permission
 	 * check.
 	 */
-	if (!bms_is_member(glob->basenodeSessionVarid, glob->checkSelectPermVarids))
-		result->exclSelectPermCheckVarid = glob->basenodeSessionVarid;
+	if (OidIsValid(root->parse->resultVariable) &&
+		!root->parse->is_Variable_assigned)
+		result->exclSelectPermCheckVarid = root->parse->resultVariable;
+	else
+		result->exclSelectPermCheckVarid = InvalidOid;
 
 	result->stmt_location = parse->stmt_location;
 	result->stmt_len = parse->stmt_len;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index 95e8538117..e627185388 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -2171,18 +2171,35 @@ fix_param_node(PlannerInfo *root, Param *p)
 		 * indirection. In this case we must not do SELECT permission check.
 		 * In other cases we have do SELECT permission check.
 		 */
-		if (p->parambasenode)
+		if (!p->parambasenode)
 		{
-			/* the variable can be used only once as base node */
-			Assert(!(OidIsValid(root->glob->basenodeSessionVarid)));
+			bool		exists = false;
+			PlannerInfo *temp_root	= NULL;
+			temp_root = root;
+			do
+			{
+				if (temp_root->parse &&
+					OidIsValid(temp_root->parse->resultVariable) &&
+					(p->paramvarid == temp_root->parse->resultVariable))
+				{
+					exists = true;
+					break;
+				}
+				temp_root = temp_root->parent_root;
+			}
+			while (temp_root != NULL);
 
-			root->glob->basenodeSessionVarid = p->paramvarid;
-		}
-		else
-		{
-			root->glob->checkSelectPermVarids =
-							  bms_add_member(root->glob->checkSelectPermVarids,
-											 p->paramvarid);
+			if (exists)
+			{
+				temp_root = root;
+				do
+				{
+					if (temp_root->parse)
+						temp_root->parse->is_Variable_assigned = true;
+					temp_root = temp_root->parent_root;
+				}
+				while (temp_root != NULL);
+			}
 		}
 
 		return (Node *) p;
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index b417b56c2a..653fa52113 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -145,6 +145,15 @@ typedef struct Query
 	/* target variable of LET statement */
 	Oid			resultVariable;
 
+	/* true if session variables is in both sides of LET statement, this flag is
+	 * set recursively for parent_root. for LET statement, when we assigned
+	 * value to a variable, we don't need check users SELECT priviledge for that
+	 * variable, but if variable is in both side of LET statement, then we do
+	 * need check the session variable SELECT priviledge again current user.
+	 * for example LET v1.a = v1.a.
+	*/
+	bool		is_Variable_assigned;
+
 	/* has aggregates in tlist or havingQual */
 	bool		hasAggs pg_node_attr(query_jumble_ignore);
 	/* has window functions in tlist */
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index f1a93d5b34..880bb2a26f 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -166,12 +166,6 @@ typedef struct PlannerGlobal
 
 	/* list of used session variables */
 	List	   *sessionVariables;
-
-	/* oid of session variable used like base node for assignment indirection */
-	Oid			basenodeSessionVarid;
-
-	/* set of session variables where execute permission check is required */
-	Bitmapset *checkSelectPermVarids;
 } PlannerGlobal;
 
 /* macro for fetching the Plan associated with a SubPlan node */
-- 
2.34.1