From 797ed6f9a628f270e3848cc5be3b6fc4e1a59194 Mon Sep 17 00:00:00 2001 From: Dipesh Dhameliya Date: Fri, 11 Apr 2025 20:05:20 +0530 Subject: [PATCH] Avoid checking permission of Babelfish temp tables on parallel worker (#569) 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 --- src/backend/executor/execMain.c | 30 ++++++++++++++++++++++++++--- src/backend/executor/execParallel.c | 19 ++++++++++++++++++ src/backend/nodes/read.c | 18 +++++++++++++++++ src/include/access/parallel.h | 4 ---- src/include/executor/execParallel.h | 11 +++++++++++ src/include/executor/executor.h | 4 ++++ src/include/nodes/nodes.h | 1 + 7 files changed, 80 insertions(+), 7 deletions(-) diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c index cab0c73a033..2f86b42fc71 100644 --- a/src/backend/executor/execMain.c +++ b/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); +} diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c index f1fd7f7e8b2..1e1d48b5119 100644 --- a/src/backend/executor/execParallel.c +++ b/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); diff --git a/src/backend/nodes/read.c b/src/backend/nodes/read.c index b42a32e8737..e29ca1aaf61 100644 --- a/src/backend/nodes/read.c +++ b/src/backend/nodes/read.c @@ -22,6 +22,7 @@ #include #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; +} \ No newline at end of file diff --git a/src/include/access/parallel.h b/src/include/access/parallel.h index 9126fd0e5da..0b4707a3103 100644 --- a/src/include/access/parallel.h +++ b/src/include/access/parallel.h @@ -82,9 +82,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; @@ -92,5 +89,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 3a1b11326f8..9f53548525d 100644 --- a/src/include/executor/execParallel.h +++ b/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 */ diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h index 4496ada1955..7dbae0ac7e9 100644 --- a/src/include/executor/executor.h +++ b/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, diff --git a/src/include/nodes/nodes.h b/src/include/nodes/nodes.h index 50e614c21c0..dbb97d007e8 100644 --- a/src/include/nodes/nodes.h +++ b/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