Avoid checking permission of Babelfish temp tables on parallel worker

src/backend/executor/execMain.c

@@ -81,6 +81,9 @@ check_rowcount_hook_type check_rowcount_hook = NULL;
 /* Hook for plugin to get control in ExecCheckRTPerms() */
 ExecutorCheckPerms_hook_type ExecutorCheckPerms_hook = NULL;
 
+/* Hook for plugin to get control in ExecCheckRTEPerms() */
+ExecCheckRTEPerms_hook_type ExecCheckRTEPerms_hook = NULL;
+
 /* decls for local routines only used within this module */
 static void InitPlan(QueryDesc *queryDesc, int eflags);
 static void CheckValidRowMarkRel(Relation rel, RowMarkType markType);
@@ -621,6 +624,19 @@ ExecCheckRTEPerms(RangeTblEntry *rte)
 
 	relOid = rte->relid;
 
+	/* 
+	 * 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.
+	 * Note - This hook must not be used outside of Babelfish parallel worker
+	 */
+	if (ExecCheckRTEPerms_hook &&
+		IsBabelfishParallelWorker() &&
+		(*ExecCheckRTEPerms_hook)(rte))
+		return true;
+
 	/*
 	 * userid to check as: current user unless we have a setuid indication.
 	 *
@@ -813,10 +829,9 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	int			i;
 
 	/*
-	 * Do permissions checks if not Babelfish parallel worker
+	 * Do permissions checks
 	 */
-	if (!IsBabelfishParallelWorker())
-		ExecCheckRTPerms(rangeTable, true);
+	ExecCheckRTPerms(rangeTable, true);
 
 	/*
 	 * initialize the node's execution state
@@ -3043,3 +3058,12 @@ EvalPlanQualEnd(EPQState *epqstate)
 	epqstate->relsubs_done = NULL;
 	epqstate->epqExtra->relsubs_blocked = NULL;
 }
+
+/*
+ * ExecCheckRTEPerms_wrapper - wrapper around ExecCheckRTEPerms
+ */
+bool
+ExecCheckRTEPerms_wrapper(RangeTblEntry *rte)
+{
+	return ExecCheckRTEPerms(rte);
+}

src/backend/executor/execParallel.c

@@ -124,6 +124,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 *node,
@@ -684,6 +687,10 @@ ExecInitParallelPlan(PlanState *planstate, EState *estate,
 						   mul_size(PARALLEL_TUPLE_QUEUE_SIZE, pcxt->nworkers));
 	shm_toc_estimate_keys(&pcxt->estimator, 1);
 
+	/* 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
 	 * count of how many PlanState nodes there are.
@@ -767,6 +774,12 @@ 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 context */
+	if (ExecInitParallelPlan_hook)
+	{
+		(*ExecInitParallelPlan_hook)(estate, pcxt,false);
+	}
+
 	/* Set up the tuple queues that the workers will write into. */
 	pei->tqueue = ExecParallelSetupTupleQueues(pcxt, false);
 
@@ -1427,6 +1440,12 @@ 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);
 
+	/* Give extension chance to retrieve additional context shared by leader node */
+	if (ParallelQueryMain_hook)
+	{
+		(*ParallelQueryMain_hook)(toc);
+	}
+
 	/* Start up the executor */
 	queryDesc->plannedstmt->jitFlags = fpes->jit_flags;
 	ExecutorStart(queryDesc, fpes->eflags);

src/backend/nodes/read.c

@@ -22,6 +22,7 @@
 #include <ctype.h>
 
 #include "common/string.h"
+#include "nodes/bitmapset.h"
 #include "nodes/pg_list.h"
 #include "nodes/readfuncs.h"
 #include "nodes/value.h"
@@ -478,3 +479,20 @@ nodeRead(const char *token, int tok_len)
 
 	return (void *) result;
 }
+
+/*
+ * Helper function for Babelfish to build Bitmapset from string.
+ */
+Bitmapset *
+stringToBms(const char *str)
+{
+	Bitmapset *ret;
+	const char *save_strtok;
+
+	save_strtok = pg_strtok_ptr;
+	pg_strtok_ptr = str;		/* point pg_strtok at the string to read */
+
+	ret = readBitmapset();
+	pg_strtok_ptr = save_strtok;
+	return ret;
+}

src/include/access/parallel.h

@@ -82,15 +82,11 @@ 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;
 
 typedef void (*bbf_ParallelWorkerMain_hook_type)(shm_toc *toc);
 extern PGDLLIMPORT bbf_ParallelWorkerMain_hook_type bbf_ParallelWorkerMain_hook;
 
-
 #endif							/* PARALLEL_H */

src/include/executor/execParallel.h

@@ -48,4 +48,15 @@ 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;
+
+/*
+ * 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;
+
 #endif							/* EXECPARALLEL_H */

src/include/executor/executor.h

@@ -99,6 +99,9 @@ extern PGDLLIMPORT ExecUpdateResultTypeTL_hook_type ExecUpdateResultTypeTL_hook;
 typedef bool (*check_rowcount_hook_type) (int es_processed);
 extern PGDLLIMPORT check_rowcount_hook_type check_rowcount_hook;
 
+typedef bool (*ExecCheckRTEPerms_hook_type) (RangeTblEntry *rte);
+extern PGDLLEXPORT ExecCheckRTEPerms_hook_type ExecCheckRTEPerms_hook;
+
 /*
  * prototypes from functions in execAmi.c
  */
@@ -212,6 +215,7 @@ extern void standard_ExecutorEnd(QueryDesc *queryDesc);
 extern void ExecutorRewind(QueryDesc *queryDesc);
 extern bool ExecCheckRTPerms(List *rangeTable, bool ereport_on_violation);
 extern void CheckValidResultRel(ResultRelInfo *resultRelInfo, CmdType operation);
+extern bool ExecCheckRTEPerms_wrapper(RangeTblEntry *rte);
 extern void InitResultRelInfo(ResultRelInfo *resultRelInfo,
 							  Relation resultRelationDesc,
 							  Index resultRelationIndex,

src/include/nodes/nodes.h

@@ -654,6 +654,7 @@ extern bool *readBoolCols(int numCols);
 extern int *readIntCols(int numCols);
 extern Oid *readOidCols(int numCols);
 extern int16 *readAttrNumberCols(int numCols);
+extern struct Bitmapset *stringToBms(const char *str);
 
 /*
  * nodes/copyfuncs.c