From 26bde8e29f2ec5921da63a01dcb30bd912e4e1c5 Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Tue, 1 Apr 2025 14:43:28 +0000
Subject: [PATCH 01/10] Dummy PR to test fix for BABEL-5703

---
 src/backend/executor/execMain.c     | 11 +++++++---
 src/backend/executor/execParallel.c | 33 +++++++++++++++++++++++++++++
 src/include/executor/executor.h     |  2 ++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 492c467698d..5e28eef4d75 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -62,6 +62,8 @@
 #include "utils/snapmgr.h"
 
 
+Bitmapset *tempRelids = NULL;
+
 /* Hooks for plugins to get control in ExecutorStart/Run/Finish/End */
 ExecutorStart_hook_type ExecutorStart_hook = NULL;
 ExecutorRun_hook_type ExecutorRun_hook = NULL;
@@ -615,6 +617,10 @@ ExecCheckPermissions(List *rangeTable, List *rteperminfos,
 		RTEPermissionInfo *perminfo = lfirst_node(RTEPermissionInfo, l);
 
 		Assert(OidIsValid(perminfo->relid));
+
+		if (IsBabelfishParallelWorker() && bms_is_member(perminfo->relid, tempRelids))
+			continue;
+
 		result = ExecCheckOneRelPerms(perminfo);
 		if (!result)
 		{
@@ -840,10 +846,9 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	int			i;
 
 	/*
-	 * Do permissions checks if not Babelfish parallel worker
+	 * Do permissions checks
 	 */
-	if (!IsBabelfishParallelWorker())
-		ExecCheckPermissions(rangeTable, plannedstmt->permInfos, true);
+	ExecCheckPermissions(rangeTable, plannedstmt->permInfos, true);
 
 	/*
 	 * initialize the node's execution state
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 8c53d1834e9..a3e1ef1e3e4 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -65,6 +65,8 @@
 #define PARALLEL_KEY_JIT_INSTRUMENTATION UINT64CONST(0xE000000000000009)
 #define PARALLEL_KEY_WAL_USAGE			UINT64CONST(0xE00000000000000A)
 
+#define PARALLEL_KEY_TEMP_RELIDS		UINT64CONST(0xE00000000000000B)
+
 #define PARALLEL_TUPLE_QUEUE_SIZE		65536
 
 /*
@@ -608,6 +610,23 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 	Size		dsa_minsize = dsa_minimum_size();
 	char	   *query_string;
 	int			query_len;
+	ListCell	*lc;
+	char	   *temp_relids_str = NULL;
+	char	   *temp_relids_space;
+	Bitmapset	*temp_relids = NULL;
+
+	/* Put this under BBF extension and have dialect check */
+	foreach(lc, estate->es_range_table)
+	{
+		RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
+		if (rte->rtekind == RTE_RELATION &&
+			OidIsValid(rte->relid) && 
+			 get_rel_persistence(rte->relid) == 't')
+		{
+			temp_relids = bms_add_member(temp_relids, rte->relid);
+		}
+	}
+	temp_relids_str = bmsToString(temp_relids);
 
 	/*
 	 * Force any initplan outputs that we're going to pass to workers to be
@@ -678,6 +697,12 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 						   mul_size(sizeof(WalUsage), pcxt->nworkers));
 	shm_toc_estimate_keys(&pcxt->estimator, 1);
 
+	/*
+	 * Babelfish extra context
+	 */
+	shm_toc_estimate_chunk(&pcxt->estimator, strlen(temp_relids_str) + 1);
+	shm_toc_estimate_keys(&pcxt->estimator, 1);
+
 	/* Estimate space for tuple queues. */
 	shm_toc_estimate_chunk(&pcxt->estimator,
 						   mul_size(PARALLEL_TUPLE_QUEUE_SIZE, pcxt->nworkers));
@@ -773,6 +798,11 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 	shm_toc_insert(pcxt->toc, PARALLEL_KEY_WAL_USAGE, walusage_space);
 	pei->wal_usage = walusage_space;
 
+	/* TODO: Move all this to extension. */
+	temp_relids_space = shm_toc_allocate(pcxt->toc, strlen(temp_relids_str) + 1);
+	memcpy(temp_relids_space, temp_relids_str, strlen(temp_relids_str) + 1);
+	shm_toc_insert(pcxt->toc, PARALLEL_KEY_TEMP_RELIDS, temp_relids_space);
+
 	/* Set up the tuple queues that the workers will write into. */
 	pei->tqueue = ExecParallelSetupTupleQueues(pcxt, false);
 
@@ -1410,6 +1440,7 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 	void	   *area_space;
 	dsa_area   *area;
 	ParallelWorkerContext pwcxt;
+	char	   *temp_relids_str;
 
 	/* Get fixed-size state. */
 	fpes = shm_toc_lookup(toc, PARALLEL_KEY_EXECUTOR_FIXED, false);
@@ -1431,6 +1462,8 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 
 	/* Attach to the dynamic shared memory area. */
 	area_space = shm_toc_lookup(toc, PARALLEL_KEY_DSA, false);
+	temp_relids_str = shm_toc_lookup(toc, PARALLEL_KEY_TEMP_RELIDS, false);
+	tempRelids = (Bitmapset *) stringToNode(temp_relids_str);
 	area = dsa_attach_in_place(area_space, seg);
 
 	/* Start up the executor */
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index fd900850da1..05cb2c41d69 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -691,4 +691,6 @@ extern ResultRelInfo *ExecLookupResultRelByOid(ModifyTableState *node,
 											   bool missing_ok,
 											   bool update_cache);
 
+extern Bitmapset *tempRelids;
+
 #endif							/* EXECUTOR_H  */

From d306a3e1f1e301e8052f0c436b56e6bcaf2fcfa8 Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Thu, 3 Apr 2025 09:05:20 +0000
Subject: [PATCH 02/10] Better separation for code using hooks

---
 src/backend/executor/execMain.c     | 10 ++++++-
 src/backend/executor/execParallel.c | 45 +++++++++--------------------
 src/include/access/parallel.h       |  4 ---
 src/include/executor/execParallel.h |  6 ++++
 src/include/executor/executor.h     |  5 ++--
 5 files changed, 31 insertions(+), 39 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 5e28eef4d75..fd85c0ccbe2 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -72,6 +72,8 @@ ExecutorEnd_hook_type ExecutorEnd_hook = NULL;
 
 TriggerRecuresiveCheck_hook_type TriggerRecuresiveCheck_hook = NULL;
 check_rowcount_hook_type check_rowcount_hook = NULL;
+skip_ExecutorCheckPerms_hook_type skip_ExecutorCheckPerms_hook = NULL;
+
 /* Hook for plugin to get control in ExecCheckPermissions() */
 ExecutorCheckPerms_hook_type ExecutorCheckPerms_hook = NULL;
 
@@ -618,7 +620,13 @@ ExecCheckPermissions(List *rangeTable, List *rteperminfos,
 
 		Assert(OidIsValid(perminfo->relid));
 
-		if (IsBabelfishParallelWorker() && bms_is_member(perminfo->relid, tempRelids))
+		/* 
+		 * Babelfish specific logic - skip permission check for temp table
+		 * under parallel worker
+		 */
+		if (IsBabelfishParallelWorker() &&
+			skip_ExecutorCheckPerms_hook &&
+			(*skip_ExecutorCheckPerms_hook)(perminfo->relid))
 			continue;
 
 		result = ExecCheckOneRelPerms(perminfo);
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index a3e1ef1e3e4..c48ad358f5f 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -65,8 +65,6 @@
 #define PARALLEL_KEY_JIT_INSTRUMENTATION UINT64CONST(0xE000000000000009)
 #define PARALLEL_KEY_WAL_USAGE			UINT64CONST(0xE00000000000000A)
 
-#define PARALLEL_KEY_TEMP_RELIDS		UINT64CONST(0xE00000000000000B)
-
 #define PARALLEL_TUPLE_QUEUE_SIZE		65536
 
 /*
@@ -124,6 +122,9 @@ typedef struct ExecParallelInitializeDSMContext
 	int			nnodes;
 } ExecParallelInitializeDSMContext;
 
+ExecInitParallelPlan_hook_type ExecInitParallelPlan_hook = NULL;
+ParallelQueryMain_hook_type ParallelQueryMain_hook = NULL;
+
 /* Helper functions that run in the parallel leader. */
 static char *ExecSerializePlan(Plan *plan, EState *estate);
 static bool ExecParallelEstimate(PlanState *planstate,
@@ -610,23 +611,6 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 	Size		dsa_minsize = dsa_minimum_size();
 	char	   *query_string;
 	int			query_len;
-	ListCell	*lc;
-	char	   *temp_relids_str = NULL;
-	char	   *temp_relids_space;
-	Bitmapset	*temp_relids = NULL;
-
-	/* Put this under BBF extension and have dialect check */
-	foreach(lc, estate->es_range_table)
-	{
-		RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
-		if (rte->rtekind == RTE_RELATION &&
-			OidIsValid(rte->relid) && 
-			 get_rel_persistence(rte->relid) == 't')
-		{
-			temp_relids = bms_add_member(temp_relids, rte->relid);
-		}
-	}
-	temp_relids_str = bmsToString(temp_relids);
 
 	/*
 	 * Force any initplan outputs that we're going to pass to workers to be
@@ -697,12 +681,6 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 						   mul_size(sizeof(WalUsage), pcxt->nworkers));
 	shm_toc_estimate_keys(&pcxt->estimator, 1);
 
-	/*
-	 * Babelfish extra context
-	 */
-	shm_toc_estimate_chunk(&pcxt->estimator, strlen(temp_relids_str) + 1);
-	shm_toc_estimate_keys(&pcxt->estimator, 1);
-
 	/* Estimate space for tuple queues. */
 	shm_toc_estimate_chunk(&pcxt->estimator,
 						   mul_size(PARALLEL_TUPLE_QUEUE_SIZE, pcxt->nworkers));
@@ -798,10 +776,11 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 	shm_toc_insert(pcxt->toc, PARALLEL_KEY_WAL_USAGE, walusage_space);
 	pei->wal_usage = walusage_space;
 
-	/* TODO: Move all this to extension. */
-	temp_relids_space = shm_toc_allocate(pcxt->toc, strlen(temp_relids_str) + 1);
-	memcpy(temp_relids_space, temp_relids_str, strlen(temp_relids_str) + 1);
-	shm_toc_insert(pcxt->toc, PARALLEL_KEY_TEMP_RELIDS, temp_relids_space);
+	/* Give extension a chance to share additional details */
+	if (ExecInitParallelPlan_hook)
+	{
+		(*ExecInitParallelPlan_hook)(estate, pcxt,false);
+	}
 
 	/* Set up the tuple queues that the workers will write into. */
 	pei->tqueue = ExecParallelSetupTupleQueues(pcxt, false);
@@ -1440,7 +1419,6 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 	void	   *area_space;
 	dsa_area   *area;
 	ParallelWorkerContext pwcxt;
-	char	   *temp_relids_str;
 
 	/* Get fixed-size state. */
 	fpes = shm_toc_lookup(toc, PARALLEL_KEY_EXECUTOR_FIXED, false);
@@ -1462,10 +1440,13 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 
 	/* Attach to the dynamic shared memory area. */
 	area_space = shm_toc_lookup(toc, PARALLEL_KEY_DSA, false);
-	temp_relids_str = shm_toc_lookup(toc, PARALLEL_KEY_TEMP_RELIDS, false);
-	tempRelids = (Bitmapset *) stringToNode(temp_relids_str);
 	area = dsa_attach_in_place(area_space, seg);
 
+	if (ParallelQueryMain_hook)
+	{
+		(*ParallelQueryMain_hook)(toc);
+	}
+
 	/* Start up the executor */
 	queryDesc->plannedstmt->jitFlags = fpes->jit_flags;
 	ExecutorStart(queryDesc, fpes->eflags);
diff --git a/src/include/access/parallel.h b/src/include/access/parallel.h
index 49a2af7050e..35f2c9874fa 100644
--- a/src/include/access/parallel.h
+++ b/src/include/access/parallel.h
@@ -81,9 +81,6 @@ extern void ParallelWorkerMain(Datum main_arg);
 /* Below helpers are added to support parallel workers in Babelfish context */
 extern bool IsBabelfishParallelWorker(void);
 
-/* Key for BabelfishFixedParallelState */
-#define BABELFISH_PARALLEL_KEY_FIXED        UINT64CONST(0xBBF0000000000001)
-
 /* Hooks for communicating babelfish related information to parallel worker */
 typedef void (*bbf_InitializeParallelDSM_hook_type)(ParallelContext *pcxt, bool estimate);
 extern PGDLLIMPORT bbf_InitializeParallelDSM_hook_type bbf_InitializeParallelDSM_hook;
@@ -91,5 +88,4 @@ extern PGDLLIMPORT bbf_InitializeParallelDSM_hook_type bbf_InitializeParallelDSM
 typedef void (*bbf_ParallelWorkerMain_hook_type)(shm_toc *toc);
 extern PGDLLIMPORT bbf_ParallelWorkerMain_hook_type bbf_ParallelWorkerMain_hook;
 
-
 #endif							/* PARALLEL_H */
diff --git a/src/include/executor/execParallel.h b/src/include/executor/execParallel.h
index 6b8c00bb0fb..0bfff70b912 100644
--- a/src/include/executor/execParallel.h
+++ b/src/include/executor/execParallel.h
@@ -48,4 +48,10 @@ extern void ExecParallelReinitialize(PlanState *planstate,
 
 extern void ParallelQueryMain(dsm_segment *seg, shm_toc *toc);
 
+typedef void (*ParallelQueryMain_hook_type)(shm_toc *toc);
+extern PGDLLIMPORT ParallelQueryMain_hook_type ParallelQueryMain_hook;
+
+typedef void (*ExecInitParallelPlan_hook_type)(EState *estate, ParallelContext *pcxt, bool estimate);
+extern PGDLLIMPORT ExecInitParallelPlan_hook_type ExecInitParallelPlan_hook;
+
 #endif							/* EXECPARALLEL_H */
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 05cb2c41d69..f21b3a601ab 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -111,6 +111,9 @@ extern PGDLLIMPORT ExecUpdateResultTypeTL_hook_type ExecUpdateResultTypeTL_hook;
 typedef bool (*check_rowcount_hook_type) (int es_processed);
 extern PGDLLEXPORT check_rowcount_hook_type check_rowcount_hook;
 
+typedef bool (*skip_ExecutorCheckPerms_hook_type) (Oid relid);
+extern PGDLLEXPORT skip_ExecutorCheckPerms_hook_type skip_ExecutorCheckPerms_hook;
+
 /*
  * prototypes from functions in execAmi.c
  */
@@ -691,6 +694,4 @@ extern ResultRelInfo *ExecLookupResultRelByOid(ModifyTableState *node,
 											   bool missing_ok,
 											   bool update_cache);
 
-extern Bitmapset *tempRelids;
-
 #endif							/* EXECUTOR_H  */

From eb9f330fef3b9ea988a1d6e5c18a1bbaba5cd0d6 Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Thu, 3 Apr 2025 09:32:53 +0000
Subject: [PATCH 03/10] Better separation for code using hooks

---
 src/backend/executor/execMain.c     | 2 --
 src/backend/executor/execParallel.c | 8 +++++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index fd85c0ccbe2..762f0114466 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -62,8 +62,6 @@
 #include "utils/snapmgr.h"
 
 
-Bitmapset *tempRelids = NULL;
-
 /* Hooks for plugins to get control in ExecutorStart/Run/Finish/End */
 ExecutorStart_hook_type ExecutorStart_hook = NULL;
 ExecutorRun_hook_type ExecutorRun_hook = NULL;
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index c48ad358f5f..6acca4921b9 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -686,6 +686,12 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 						   mul_size(PARALLEL_TUPLE_QUEUE_SIZE, pcxt->nworkers));
 	shm_toc_estimate_keys(&pcxt->estimator, 1);
 
+	/* Give extension a chance to share additional details */
+	if (ExecInitParallelPlan_hook)
+	{
+		(*ExecInitParallelPlan_hook)(estate, pcxt, true);
+	}
+
 	/*
 	 * Give parallel-aware nodes a chance to add to the estimates, and get a
 	 * count of how many PlanState nodes there are.
@@ -1442,7 +1448,7 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 	area_space = shm_toc_lookup(toc, PARALLEL_KEY_DSA, false);
 	area = dsa_attach_in_place(area_space, seg);
 
-	if (ParallelQueryMain_hook)
+	if (IsBabelfishParallelWorker() && ParallelQueryMain_hook)
 	{
 		(*ParallelQueryMain_hook)(toc);
 	}

From 6f8114576f37b6ab2680295aa3231dd876d34d66 Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Thu, 3 Apr 2025 14:29:08 +0000
Subject: [PATCH 04/10] expose function to let extension do perm checks

---
 src/backend/executor/execMain.c | 3 +--
 src/include/executor/executor.h | 1 +
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 762f0114466..b162576014d 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -86,7 +86,6 @@ static void ExecutePlan(QueryDesc *queryDesc,
 						uint64 numberTuples,
 						ScanDirection direction,
 						DestReceiver *dest);
-static bool ExecCheckOneRelPerms(RTEPermissionInfo *perminfo);
 static bool ExecCheckPermissionsModified(Oid relOid, Oid userid,
 										 Bitmapset *modifiedCols,
 										 AclMode requiredPerms);
@@ -648,7 +647,7 @@ ExecCheckPermissions(List *rangeTable, List *rteperminfos,
  * ExecCheckOneRelPerms
  *		Check access permissions for a single relation.
  */
-static bool
+bool
 ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 {
 	AclMode		requiredPerms;
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index f21b3a601ab..e7d7f0972e2 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -227,6 +227,7 @@ extern void standard_ExecutorEnd(QueryDesc *queryDesc);
 extern void ExecutorRewind(QueryDesc *queryDesc);
 extern bool ExecCheckPermissions(List *rangeTable,
 								 List *rteperminfos, bool ereport_on_violation);
+extern bool ExecCheckOneRelPerms(RTEPermissionInfo *perminfo);
 extern void CheckValidResultRel(ResultRelInfo *resultRelInfo, CmdType operation,
 								List *mergeActions);
 extern void InitResultRelInfo(ResultRelInfo *resultRelInfo,

From 0010f94d8b2a3e78d492f52603b589e7af7441dc Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Fri, 4 Apr 2025 05:48:47 +0000
Subject: [PATCH 05/10] Move permission skip logic to more granular level

---
 src/backend/executor/execMain.c | 35 ++++++++++++++++++++++-----------
 src/include/executor/executor.h |  7 ++++---
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index b162576014d..60e03ab1c1d 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -70,7 +70,7 @@ ExecutorEnd_hook_type ExecutorEnd_hook = NULL;
 
 TriggerRecuresiveCheck_hook_type TriggerRecuresiveCheck_hook = NULL;
 check_rowcount_hook_type check_rowcount_hook = NULL;
-skip_ExecutorCheckPerms_hook_type skip_ExecutorCheckPerms_hook = NULL;
+ExecCheckOneRelPerms_hook_type ExecCheckOneRelPerms_hook = NULL;
 
 /* Hook for plugin to get control in ExecCheckPermissions() */
 ExecutorCheckPerms_hook_type ExecutorCheckPerms_hook = NULL;
@@ -86,6 +86,7 @@ static void ExecutePlan(QueryDesc *queryDesc,
 						uint64 numberTuples,
 						ScanDirection direction,
 						DestReceiver *dest);
+static bool ExecCheckOneRelPerms(RTEPermissionInfo *perminfo);
 static bool ExecCheckPermissionsModified(Oid relOid, Oid userid,
 										 Bitmapset *modifiedCols,
 										 AclMode requiredPerms);
@@ -616,16 +617,6 @@ ExecCheckPermissions(List *rangeTable, List *rteperminfos,
 		RTEPermissionInfo *perminfo = lfirst_node(RTEPermissionInfo, l);
 
 		Assert(OidIsValid(perminfo->relid));
-
-		/* 
-		 * Babelfish specific logic - skip permission check for temp table
-		 * under parallel worker
-		 */
-		if (IsBabelfishParallelWorker() &&
-			skip_ExecutorCheckPerms_hook &&
-			(*skip_ExecutorCheckPerms_hook)(perminfo->relid))
-			continue;
-
 		result = ExecCheckOneRelPerms(perminfo);
 		if (!result)
 		{
@@ -647,7 +638,7 @@ ExecCheckPermissions(List *rangeTable, List *rteperminfos,
  * ExecCheckOneRelPerms
  *		Check access permissions for a single relation.
  */
-bool
+static bool
 ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 {
 	AclMode		requiredPerms;
@@ -659,6 +650,17 @@ ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 	requiredPerms = perminfo->requiredPerms;
 	Assert(requiredPerms != 0);
 
+	/* 
+	 * Some extension may want to avoid permission check for some special
+	 * cases. For example, Babelfish temp table is implemented using ENR
+	 * which is not shared with parallel worker and parallel operations
+	 * are not allowed for temp table in Postgres. So give extension a
+	 * chance to skip permission check for such use cases.
+	 */
+	if (ExecCheckOneRelPerms_hook &&
+		(*ExecCheckOneRelPerms_hook)(perminfo))
+		return true;
+
 	/*
 	 * userid to check as: current user unless we have a setuid indication.
 	 *
@@ -3056,3 +3058,12 @@ EvalPlanQualEnd(EPQState *epqstate)
 	epqstate->relsubs_done = NULL;
 	epqstate->relsubs_blocked = NULL;
 }
+
+/*
+ * ExecCheckOneRelPerms_wrapper - wrapper around ExecCheckOneRelPerms
+ */
+bool
+ExecCheckOneRelPerms_wrapper(RTEPermissionInfo *perminfo)
+{
+	return ExecCheckOneRelPerms(perminfo);
+}
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index e7d7f0972e2..047d3e9d9a2 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -111,8 +111,8 @@ extern PGDLLIMPORT ExecUpdateResultTypeTL_hook_type ExecUpdateResultTypeTL_hook;
 typedef bool (*check_rowcount_hook_type) (int es_processed);
 extern PGDLLEXPORT check_rowcount_hook_type check_rowcount_hook;
 
-typedef bool (*skip_ExecutorCheckPerms_hook_type) (Oid relid);
-extern PGDLLEXPORT skip_ExecutorCheckPerms_hook_type skip_ExecutorCheckPerms_hook;
+typedef bool (*ExecCheckOneRelPerms_hook_type) (RTEPermissionInfo *perminfo);
+extern PGDLLEXPORT ExecCheckOneRelPerms_hook_type ExecCheckOneRelPerms_hook;
 
 /*
  * prototypes from functions in execAmi.c
@@ -227,7 +227,8 @@ extern void standard_ExecutorEnd(QueryDesc *queryDesc);
 extern void ExecutorRewind(QueryDesc *queryDesc);
 extern bool ExecCheckPermissions(List *rangeTable,
 								 List *rteperminfos, bool ereport_on_violation);
-extern bool ExecCheckOneRelPerms(RTEPermissionInfo *perminfo);
+/* Wrapper around ExecCheckOneRelPerms */
+extern bool ExecCheckOneRelPerms_wrapper(RTEPermissionInfo *perminfo);
 extern void CheckValidResultRel(ResultRelInfo *resultRelInfo, CmdType operation,
 								List *mergeActions);
 extern void InitResultRelInfo(ResultRelInfo *resultRelInfo,

From f4a4cdd66004a6a41ee044866d33d195d4feb86b Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Fri, 4 Apr 2025 06:01:54 +0000
Subject: [PATCH 06/10] More comments

---
 src/backend/executor/execParallel.c | 7 ++++---
 src/include/executor/execParallel.h | 5 +++++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 6acca4921b9..d94e9189001 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -686,7 +686,7 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 						   mul_size(PARALLEL_TUPLE_QUEUE_SIZE, pcxt->nworkers));
 	shm_toc_estimate_keys(&pcxt->estimator, 1);
 
-	/* Give extension a chance to share additional details */
+	/* Let extension estimate a dynamic shared memory needed to communicate additional context */
 	if (ExecInitParallelPlan_hook)
 	{
 		(*ExecInitParallelPlan_hook)(estate, pcxt, true);
@@ -782,7 +782,7 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 	shm_toc_insert(pcxt->toc, PARALLEL_KEY_WAL_USAGE, walusage_space);
 	pei->wal_usage = walusage_space;
 
-	/* Give extension a chance to share additional details */
+	/* Give extension a chance to share additional context */
 	if (ExecInitParallelPlan_hook)
 	{
 		(*ExecInitParallelPlan_hook)(estate, pcxt,false);
@@ -1448,7 +1448,8 @@ ParallelQueryMain(dsm_segment *seg, shm_toc *toc)
 	area_space = shm_toc_lookup(toc, PARALLEL_KEY_DSA, false);
 	area = dsa_attach_in_place(area_space, seg);
 
-	if (IsBabelfishParallelWorker() && ParallelQueryMain_hook)
+	/* Give extension chance to retrieve additional context shared by leader node */
+	if (ParallelQueryMain_hook)
 	{
 		(*ParallelQueryMain_hook)(toc);
 	}
diff --git a/src/include/executor/execParallel.h b/src/include/executor/execParallel.h
index 0bfff70b912..3c8562f95b2 100644
--- a/src/include/executor/execParallel.h
+++ b/src/include/executor/execParallel.h
@@ -51,6 +51,11 @@ extern void ParallelQueryMain(dsm_segment *seg, shm_toc *toc);
 typedef void (*ParallelQueryMain_hook_type)(shm_toc *toc);
 extern PGDLLIMPORT ParallelQueryMain_hook_type ParallelQueryMain_hook;
 
+/*
+ * When estimate = true passed then caller wants extension to estimate a dynamic shared memory (DSM)
+ * needed by that extension to communicate additional context with Parallel worker.
+ * When estimate = false then caller wants to insert additional context to DSM.
+ */
 typedef void (*ExecInitParallelPlan_hook_type)(EState *estate, ParallelContext *pcxt, bool estimate);
 extern PGDLLIMPORT ExecInitParallelPlan_hook_type ExecInitParallelPlan_hook;
 

From 0b3cf1487c91fa8e3336601ba038ded97eb6e673 Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Fri, 4 Apr 2025 14:28:53 +0000
Subject: [PATCH 07/10] check babelfish parallel worker before skipping perm
 check

---
 src/backend/executor/execMain.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 60e03ab1c1d..34127379e35 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -657,7 +657,8 @@ ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 	 * are not allowed for temp table in Postgres. So give extension a
 	 * chance to skip permission check for such use cases.
 	 */
-	if (ExecCheckOneRelPerms_hook &&
+	if (IsBabelfishParallelWorker() &&
+		ExecCheckOneRelPerms_hook &&
 		(*ExecCheckOneRelPerms_hook)(perminfo))
 		return true;
 

From c074068ba2fc6d9333b04e9c2ffdcf364ad565db Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Tue, 8 Apr 2025 12:33:58 +0000
Subject: [PATCH 08/10] Better comment

---
 src/backend/executor/execMain.c | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 34127379e35..6d1b58d62b0 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -651,11 +651,11 @@ ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 	Assert(requiredPerms != 0);
 
 	/* 
-	 * Some extension may want to avoid permission check for some special
-	 * cases. For example, Babelfish temp table is implemented using ENR
-	 * which is not shared with parallel worker and parallel operations
-	 * are not allowed for temp table in Postgres. So give extension a
-	 * chance to skip permission check for such use cases.
+	 * Babelfish specific logic - Babelfish temp table is implemented
+	 * using ENR which is not shared with parallel worker and parallel
+	 * operations are not allowed for temp table in Postgres. Babelfish
+	 * can skip permission check for such use cases under parallel worker
+	 * using this hook.
 	 */
 	if (IsBabelfishParallelWorker() &&
 		ExecCheckOneRelPerms_hook &&

From 14e63fbf60b1912358b5257adcdb0d7fb89245fc Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Wed, 9 Apr 2025 03:50:54 +0000
Subject: [PATCH 09/10] Better comments

---
 src/backend/executor/execMain.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 6d1b58d62b0..81f59adbb19 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -656,6 +656,7 @@ ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 	 * operations are not allowed for temp table in Postgres. Babelfish
 	 * can skip permission check for such use cases under parallel worker
 	 * using this hook.
+	 * Note - This hook must not be used outside of Babelfish parallel worker
 	 */
 	if (IsBabelfishParallelWorker() &&
 		ExecCheckOneRelPerms_hook &&

From 4132d4cf65a3f1ce039b72e2a82264bd5d1815fc Mon Sep 17 00:00:00 2001
From: Dipesh Dhameliya <dddhamel@amazon.com>
Date: Wed, 9 Apr 2025 13:32:14 +0000
Subject: [PATCH 10/10] Adress comments

---
 src/backend/executor/execMain.c     | 4 ++--
 src/backend/executor/execParallel.c | 2 --
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 81f59adbb19..9369d1265e2 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -658,8 +658,8 @@ ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
 	 * using this hook.
 	 * Note - This hook must not be used outside of Babelfish parallel worker
 	 */
-	if (IsBabelfishParallelWorker() &&
-		ExecCheckOneRelPerms_hook &&
+	if (ExecCheckOneRelPerms_hook &&
+		IsBabelfishParallelWorker() &&
 		(*ExecCheckOneRelPerms_hook)(perminfo))
 		return true;
 
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index d94e9189001..e6731f9505c 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -688,9 +688,7 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 
 	/* Let extension estimate a dynamic shared memory needed to communicate additional context */
 	if (ExecInitParallelPlan_hook)
-	{
 		(*ExecInitParallelPlan_hook)(estate, pcxt, true);
-	}
 
 	/*
 	 * Give parallel-aware nodes a chance to add to the estimates, and get a