Avoid checking permission of Babelfish temp tables on parallel worker (#569) (#571)

Consider following facts,

Babelfish temp tables are implemented using ENR which is not shared between different backends. So if Parallel worker tries to check permissions on Babelfish then it will fail.
Any user should be able to access Babelfish temp tables under given session.
Postgres by default does not allow parallel operations on temp tables. Attempt to do so will result in run time error.
Due to above facts, we should avoid permission check on temp tables within parallel workers while ensuring that leader does required permission check on other tables. This commits achieves this behaviour by introducing following three hooks,

ParallelQueryMain_hook -- Hook that allows other extensions to pass on additional details from Leader node to parallel worker. For example, Babelfish extension can pass details of Babelfish temp table defined under current session with Parallel workers.

ExecInitParallelPlan_hook -- Hook that allows Parallel worker to gather additional details passed by Leader node. For example, Babelfish extension can collect the details of Babelfish temp table shared by Leader node so that it can avoid permission checks.

ExecCheckRTEPerms_hook -- Hook that allows extension control permission checking on given relation/table. For example, Babelfish can use it to avoid permission check on temp tables under parallel worker.

Also, this changes introduces stringToBms(...) API using which we can directly build Bitmapset from given string. This is
especially useful by Babelfish parallel worker to prepare Bitmapset of temp relids from string shared by Leader node.

Task: BABEL-5703
Signed-off-by: Dipesh Dhameliya <[email protected]>

Author: Deepesh125

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