Skip to content

Commit 7b60037

Browse files
Deepesh125Dipesh Dhameliya
Deepesh125
authored and
Dipesh Dhameliya
committed
Avoid checking permission of Babelfish temp tables on parallel worker (babelfish-for-postgresql#3644)
Consider following facts, 1. 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. 2. Any user should be able to access Babelfish temp tables under given session. 3. 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 implementating following three hooks, bbf_ParallelQueryMain_hook -- implements ParallelQueryMain_hook. It constructs Bitmapset (list of temp table oids) by looping through es_range_table checking if given relation is temp table, checks permission if it is and add it to the set. It will then pass this set of temp table defined under current with Parallel workers. bbf_ExecInitParallelPlan -- implements ExecInitParallelPlan_hook. Parallel workers will gather additional details (oid set of temp table) passed by Leader node. This set will later be used to avoid checking permissions on temp tables. bbf_ExecCheckOneRelPerms_hook -- implements ExecCheckOneRelPerms_hook. It will avoid permission check relation if relid is part of the set of oids constructed using bbf_ExecInitParallelPlan. Task: BABEL-5703 Signed-off-by: Dipesh Dhameliya <dddhamel@amazon.com> (cherry picked from commit 6da7ff0)
1 parent 9f676e9 commit 7b60037

File tree

8 files changed

+1381
-0
lines changed

8 files changed

+1381
-0
lines changed

‎contrib/babelfishpg_tsql/Makefile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ OBJS += src/extendedproperty.o
7777
OBJS += src/fts.o
7878
OBJS += src/fts_parser.o
7979
OBJS += src/pltsql_partition.o
80+
OBJS += src/bbf_parallel_query.o
8081

8182
export ANTLR4_JAVA_BIN=java
8283
export ANTLR4_RUNTIME_LIB=-lantlr4-runtime

‎contrib/babelfishpg_tsql/src/bbf_parallel_query.c‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
2+
#include "postgres.h"
3+
4+
#include "access/parallel.h"
5+
#include "catalog/pg_class.h"
6+
#include "executor/executor.h"
7+
#include "fmgr.h"
8+
#include "nodes/bitmapset.h"
9+
#include "nodes/execnodes.h"
10+
#include "nodes/nodes.h"
11+
#include "nodes/parsenodes.h"
12+
#include "nodes/pg_list.h"
13+
#include "miscadmin.h"
14+
#include "parser/parse_relation.h"
15+
#include "parser/parser.h"
16+
#include "storage/shm_toc.h"
17+
#include "utils/acl.h"
18+
#include "utils/elog.h"
19+
#include "utils/lsyscache.h"
20+
#include "utils/queryenvironment.h"
21+
22+
#include "bbf_parallel_query.h"
23+
#include "pltsql.h"
24+
25+
/*
26+
* temp_relids - maintains relid list of temp table shared by leader node. This should be
27+
* strictly accessed within parallel worker context.
28+
*/
29+
static Bitmapset *temp_relids = NULL;
30+
31+
/*
32+
* string representation of list of temp table oids computed by Leader node to be shared
33+
* with parallel worker.
34+
*/
35+
static char *temp_relids_str = NULL;
36+
37+
/*
38+
* In Babelfish, Any user under given session should be able access the temp tables.
39+
* And Postgres doesn't allow parallel scan on temp tables. But it still does the permission
40+
* check on temp tables under parallel workers. But since Babelfish temp tables implemented
41+
* using ENR, it can't be accessed inside Parallel worker.
42+
*
43+
* So we would like to avoid permission checking on temp tables under parallel workers while
44+
* ensuring that Leader node does require permission checks on temp table.
45+
*
46+
* In order to achieve this, we (probably re)do permission checks on temp tables and share
47+
* the list of oids with parallel worker. And parallel worker will avoid permission check
48+
* on these oids.
49+
*/
50+
51+
52+
/*
53+
* bbf_ExecInitParallelPlan -- implements ExecInitParallelPlan_hook.
54+
* It iterates through es_range_tables checking persistence of given relation. Probably,
55+
* re-do permission checking (better to redo perm checking instead of never doing it) if
56+
* relation is temp table and adds oid/relid to the set. This set will be shared with
57+
* parallel worker so that parallel worker avoids permission check on temp tables.
58+
* When estimate = true passed then caller wants to estimate a dynamic shared memory (DSM)
59+
* needed by this extension to communicate additional context.
60+
* When estimate = false then caller wants to insert additional context to DSM.
61+
*/
62+
void
63+
bbf_ExecInitParallelPlan(EState *estate, ParallelContext *pcxt, bool estimate)
64+
{
65+
if (prev_ExecInitParallelPlan_hook)
66+
(*prev_ExecInitParallelPlan_hook)(estate, pcxt, estimate);
67+
68+
/*
69+
* Dialect check is not sufficient because parallel worker might be needed while
70+
* doing plpgsql function scan on leader node.
71+
*/
72+
if (!IS_TDS_CONN())
73+
return;
74+
75+
if (estimate)
76+
{
77+
ListCell *lc;
78+
Bitmapset *temp_relids_local = NULL;
79+
temp_relids_str = NULL;
80+
81+
foreach(lc, estate->es_range_table)
82+
{
83+
RTEPermissionInfo *perminfo;
84+
RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
85+
if (rte->rtekind == RTE_RELATION &&
86+
OidIsValid(rte->relid) &&
87+
get_rel_persistence(rte->relid) == RELPERSISTENCE_TEMP)
88+
{
89+
if (estate->es_plannedstmt == NULL ||
90+
list_length(estate->es_plannedstmt->permInfos) == 0)
91+
{
92+
/* If there is RTE present then corresponding permInfo must be there */
93+
ereport(ERROR,
94+
(errcode(ERRCODE_INTERNAL_ERROR),
95+
errmsg("Failed to check permission of the temporary table because " \
96+
"corresponding Perminfo is not found while launching parallel worker(s)")));
97+
}
98+
99+
/* getRTEPermissionInfo would return valid perfInfo, error will be raised otherwise */
100+
perminfo = getRTEPermissionInfo(estate->es_plannedstmt->permInfos, rte);
101+
/* probably (re)do perm check */
102+
if (!ExecCheckOneRelPerms_wrapper(perminfo))
103+
{
104+
aclcheck_error(ACLCHECK_NO_PRIV,
105+
get_relkind_objtype(get_rel_relkind(perminfo->relid)),
106+
get_rel_name(perminfo->relid));
107+
}
108+
temp_relids_local = bms_add_member(temp_relids_local, rte->relid);
109+
}
110+
}
111+
temp_relids_str = bmsToString(temp_relids_local);
112+
113+
/*
114+
* Estimate extra context for Babelfish
115+
*/
116+
shm_toc_estimate_chunk(&pcxt->estimator, strlen(temp_relids_str) + 1);
117+
shm_toc_estimate_keys(&pcxt->estimator, 1);
118+
}
119+
else
120+
{
121+
char *temp_relids_space;
122+
123+
/*temp_relids_str will never be NULL even if there is no temp tables in the query. */
124+
if (temp_relids_str == NULL)
125+
{
126+
ereport(ERROR,
127+
(errcode(ERRCODE_INTERNAL_ERROR),
128+
errmsg("Unexpected list of temp table relids")));
129+
}
130+
131+
temp_relids_space = shm_toc_allocate(pcxt->toc, strlen(temp_relids_str) + 1);
132+
memcpy(temp_relids_space, temp_relids_str, strlen(temp_relids_str) + 1);
133+
shm_toc_insert(pcxt->toc, BABELFISH_PARALLEL_KEY_TEMP_RELIDS, temp_relids_space);
134+
135+
/* And reset temp_relids_str */
136+
pfree(temp_relids_str);
137+
temp_relids_str = NULL;
138+
}
139+
}
140+
/*
141+
* bbf_ParallelQueryMain -- implements ParallelQueryMain_hook.
142+
* It constructs temp_relids which represents oid list of temp tables communicated by Leader node.
143+
* warning: should stricktly call under parallel worker.
144+
*/
145+
void
146+
bbf_ParallelQueryMain(shm_toc *toc)
147+
{
148+
if (prev_ParallelQueryMain_hook)
149+
(*prev_ParallelQueryMain_hook)(toc);
150+
151+
/* Another line of defense to make sure no regular backend calls this function. */
152+
if (!IsBabelfishParallelWorker())
153+
{
154+
return;
155+
}
156+
157+
temp_relids = (Bitmapset *) stringToNode(shm_toc_lookup(toc,
158+
BABELFISH_PARALLEL_KEY_TEMP_RELIDS,
159+
false));
160+
}
161+
162+
/*
163+
* bbf_ExecCheckOneRelPerms -- implements ExecCheckOneRelPerms_hook.
164+
* Returns true if this is Babelfish parallel worker and provided relid is Babelfish temp table.
165+
* Note that temp_relids must have communicated by Leader node.
166+
*/
167+
bool
168+
bbf_ExecCheckOneRelPerms(RTEPermissionInfo *perminfo)
169+
{
170+
if (!OidIsValid(perminfo->relid))
171+
{
172+
ereport(ERROR,
173+
(errcode(ERRCODE_INTERNAL_ERROR),
174+
errmsg("Invalid Oid is found while checking permission of the relation")));
175+
}
176+
177+
/* Let regular permission check happen if its not Babelfish parallel worker. */
178+
if (!IsBabelfishParallelWorker())
179+
{
180+
return false;
181+
}
182+
183+
return bms_is_member(perminfo->relid, temp_relids);
184+
}
185+

‎contrib/babelfishpg_tsql/src/bbf_parallel_query.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
#include "executor/executor.h"
2+
#include "executor/execParallel.h"
3+
#include "fmgr.h"
4+
#include "nodes/execnodes.h"
5+
6+
/* Key for sharing additional context for Babelfish */
7+
#define BABELFISH_PARALLEL_KEY_FIXED UINT64CONST(0xBBF0000000000001)
8+
#define BABELFISH_PARALLEL_KEY_TEMP_RELIDS UINT64CONST(0xBBF0000000000002)
9+
10+
extern ExecInitParallelPlan_hook_type prev_ExecInitParallelPlan_hook;
11+
extern ParallelQueryMain_hook_type prev_ParallelQueryMain_hook;
12+
13+
extern void bbf_ExecInitParallelPlan(EState *estate, ParallelContext *pcxt, bool estimate);
14+
extern void bbf_ParallelQueryMain(shm_toc *toc);
15+
extern bool bbf_ExecCheckOneRelPerms(RTEPermissionInfo *perminfo);

‎contrib/babelfishpg_tsql/src/hooks.c‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@
8787
#include "multidb.h"
8888
#include "tsql_analyze.h"
8989
#include "table_variable_mvcc.h"
90+
#include "bbf_parallel_query.h"
9091

9192
#define TDS_NUMERIC_MAX_PRECISION 38
9293
extern bool babelfish_dump_restore;
@@ -323,6 +324,8 @@ static ExecFuncProc_AclCheck_hook_type prev_ExecFuncProc_AclCheck_hook = NULL;
323324
static bbf_execute_grantstmt_as_dbsecadmin_hook_type prev_bbf_execute_grantstmt_as_dbsecadmin_hook = NULL;
324325
static bbf_check_member_has_direct_priv_to_grant_role_hook_type prev_bbf_check_member_has_direct_priv_to_grant_role_hook = NULL;
325326
static validateCachedPlanSearchPath_hook_type prev_validateCachedPlanSearchPath_hook = NULL;
327+
ExecInitParallelPlan_hook_type prev_ExecInitParallelPlan_hook = NULL;
328+
ParallelQueryMain_hook_type prev_ParallelQueryMain_hook = NULL;
326329

327330
/*****************************************
328331
* Install / Uninstall
@@ -565,6 +568,14 @@ InstallExtendedHooks(void)
565568
prev_validateCachedPlanSearchPath_hook = validateCachedPlanSearchPath_hook;
566569
validateCachedPlanSearchPath_hook = pltsql_validateCachedPlanSearchPath;
567570
is_bbf_tds_connection_hook = is_bbf_tds_connection;
571+
572+
prev_ParallelQueryMain_hook = ParallelQueryMain_hook;
573+
ParallelQueryMain_hook = bbf_ParallelQueryMain;
574+
575+
prev_ExecInitParallelPlan_hook = ExecInitParallelPlan_hook;
576+
ExecInitParallelPlan_hook = bbf_ExecInitParallelPlan;
577+
578+
ExecCheckOneRelPerms_hook = bbf_ExecCheckOneRelPerms;
568579
}
569580

570581
void
@@ -648,6 +659,9 @@ UninstallExtendedHooks(void)
648659
pltsql_get_object_identity_event_trigger_hook = NULL;
649660
pltsql_allow_storing_init_privs_hook = NULL;
650661
is_bbf_tds_connection_hook = NULL;
662+
ParallelQueryMain_hook = prev_ParallelQueryMain_hook;
663+
ExecInitParallelPlan_hook = prev_ExecInitParallelPlan_hook;
664+
ExecCheckOneRelPerms_hook = NULL;
651665
}
652666

653667
/*****************************************

‎contrib/babelfishpg_tsql/src/session.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "guc.h"
2121
#include "storage/shm_toc.h"
2222
#include "collation.h"
23+
#include "bbf_parallel_query.h"
2324

2425
/* Core Session Properties */
2526

0 commit comments

Comments
 (0)