Skip to content

Commit d741f37

Browse files
Deepesh125Deepesh125
authored
Avoid checking permission of Babelfish temp tables on parallel worker (#3674)
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 implementing 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_ ExecCheckRTEPerms_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 3ddff85)
1 parent b48cd51 commit d741f37

File tree

9 files changed

+1399
-1
lines changed

9 files changed

+1399
-1
lines changed

‎contrib/babelfishpg_tsql/Makefile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ OBJS += src/tsql_for/forxml.o
7272
OBJS += src/tsql_for/forxml_old.o
7373
OBJS += src/tsql_analyze.o
7474
OBJS += src/table_variable_mvcc.o
75+
OBJS += src/bbf_parallel_query.o
7576

7677
export ANTLR4_JAVA_BIN=java
7778
export ANTLR4_RUNTIME_LIB=-lantlr4-runtime

‎contrib/babelfishpg_tsql/src/bbf_parallel_query.c‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
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+
* bbf_ExecInitParallelPlan -- implements ExecInitParallelPlan_hook.
53+
* It iterates through es_range_tables checking persistence of given relation. Probably,
54+
* re-do permission checking (better to redo perm checking instead of never doing it) if
55+
* relation is temp table and adds oid/relid to the set. This set will be shared with
56+
* parallel worker so that parallel worker avoids permission check on temp tables.
57+
* When estimate = true passed then caller wants to estimate a dynamic shared memory (DSM)
58+
* needed by this extension to communicate additional context.
59+
* When estimate = false then caller wants to insert additional context to DSM.
60+
*/
61+
void
62+
bbf_ExecInitParallelPlan(EState *estate, ParallelContext *pcxt, bool estimate)
63+
{
64+
if (prev_ExecInitParallelPlan_hook)
65+
(*prev_ExecInitParallelPlan_hook)(estate, pcxt, estimate);
66+
67+
/*
68+
* Dialect check is not sufficient because parallel worker might be needed while
69+
* doing plpgsql function scan on leader node.
70+
*/
71+
if (!IS_TDS_CONN())
72+
return;
73+
74+
if (estimate)
75+
{
76+
ListCell *lc;
77+
Bitmapset *temp_relids_local = NULL;
78+
temp_relids_str = NULL;
79+
80+
foreach(lc, estate->es_range_table)
81+
{
82+
RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
83+
if (rte->rtekind == RTE_RELATION &&
84+
OidIsValid(rte->relid) &&
85+
get_rel_persistence(rte->relid) == RELPERSISTENCE_TEMP)
86+
{
87+
/* probably (re)do perm check */
88+
if (!ExecCheckRTEPerms_wrapper(rte))
89+
{
90+
aclcheck_error(ACLCHECK_NO_PRIV,
91+
get_relkind_objtype(get_rel_relkind(rte->relid)),
92+
get_rel_name(rte->relid));
93+
}
94+
temp_relids_local = bms_add_member(temp_relids_local, rte->relid);
95+
}
96+
}
97+
temp_relids_str = bmsToString(temp_relids_local);
98+
99+
/*
100+
* Estimate extra context for Babelfish
101+
*/
102+
shm_toc_estimate_chunk(&pcxt->estimator, strlen(temp_relids_str) + 1);
103+
shm_toc_estimate_keys(&pcxt->estimator, 1);
104+
}
105+
else
106+
{
107+
char *temp_relids_space;
108+
109+
/*temp_relids_str will never be NULL even if there is no temp tables in the query. */
110+
if (temp_relids_str == NULL)
111+
{
112+
ereport(ERROR,
113+
(errcode(ERRCODE_INTERNAL_ERROR),
114+
errmsg("Unexpected list of temp table relids")));
115+
}
116+
117+
temp_relids_space = shm_toc_allocate(pcxt->toc, strlen(temp_relids_str) + 1);
118+
memcpy(temp_relids_space, temp_relids_str, strlen(temp_relids_str) + 1);
119+
shm_toc_insert(pcxt->toc, BABELFISH_PARALLEL_KEY_TEMP_RELIDS, temp_relids_space);
120+
121+
/* And reset temp_relids_str */
122+
pfree(temp_relids_str);
123+
temp_relids_str = NULL;
124+
}
125+
}
126+
/*
127+
* bbf_ParallelQueryMain -- implements ParallelQueryMain_hook.
128+
* It constructs temp_relids which represents oid list of temp tables communicated by Leader node.
129+
* warning: should stricktly call under parallel worker.
130+
*/
131+
void
132+
bbf_ParallelQueryMain(shm_toc *toc)
133+
{
134+
if (prev_ParallelQueryMain_hook)
135+
(*prev_ParallelQueryMain_hook)(toc);
136+
137+
/* Another line of defense to make sure no regular backend calls this function. */
138+
if (!IsBabelfishParallelWorker())
139+
{
140+
return;
141+
}
142+
143+
temp_relids = stringToBms(shm_toc_lookup(toc,
144+
BABELFISH_PARALLEL_KEY_TEMP_RELIDS,
145+
false));
146+
}
147+
148+
/*
149+
* bbf_ExecCheckRTEPerms -- implements ExecCheckRTEPerms_hook.
150+
* Returns true if this is Babelfish parallel worker and provided relid is Babelfish temp table.
151+
* Note that temp_relids must have communicated by Leader node.
152+
*/
153+
bool
154+
bbf_ExecCheckRTEPerms(RangeTblEntry *rte)
155+
{
156+
if (!OidIsValid(rte->relid))
157+
{
158+
ereport(ERROR,
159+
(errcode(ERRCODE_INTERNAL_ERROR),
160+
errmsg("Invalid Oid is found while checking permission of the relation")));
161+
}
162+
163+
/* Let regular permission check happen if its not Babelfish parallel worker. */
164+
if (!IsBabelfishParallelWorker())
165+
{
166+
return false;
167+
}
168+
169+
return bms_is_member(rte->relid, temp_relids);
170+
}
171+

‎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_ExecCheckRTEPerms(RangeTblEntry *rte);

‎contrib/babelfishpg_tsql/src/hooks.c‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
#include "multidb.h"
6363
#include "tsql_analyze.h"
6464
#include "table_variable_mvcc.h"
65+
#include "bbf_parallel_query.h"
6566

6667
#define TDS_NUMERIC_MAX_PRECISION 38
6768
extern bool babelfish_dump_restore;
@@ -192,7 +193,8 @@ static table_variable_satisfies_vacuum_hook_type prev_table_variable_satisfies_v
192193
static table_variable_satisfies_vacuum_horizon_hook_type prev_table_variable_satisfies_vacuum_horizon = NULL;
193194
static sortby_nulls_hook_type prev_sortby_nulls_hook = NULL;
194195
static drop_relation_refcnt_hook_type prev_drop_relation_refcnt_hook = NULL;
195-
196+
ExecInitParallelPlan_hook_type prev_ExecInitParallelPlan_hook = NULL;
197+
ParallelQueryMain_hook_type prev_ParallelQueryMain_hook = NULL;
196198

197199
/*****************************************
198200
* Install / Uninstall
@@ -327,6 +329,14 @@ InstallExtendedHooks(void)
327329

328330
bbf_InitializeParallelDSM_hook = babelfixedparallelstate_insert;
329331
bbf_ParallelWorkerMain_hook = babelfixedparallelstate_restore;
332+
333+
prev_ParallelQueryMain_hook = ParallelQueryMain_hook;
334+
ParallelQueryMain_hook = bbf_ParallelQueryMain;
335+
336+
prev_ExecInitParallelPlan_hook = ExecInitParallelPlan_hook;
337+
ExecInitParallelPlan_hook = bbf_ExecInitParallelPlan;
338+
339+
ExecCheckRTEPerms_hook = bbf_ExecCheckRTEPerms;
330340
}
331341

332342
void
@@ -381,6 +391,9 @@ UninstallExtendedHooks(void)
381391

382392
bbf_InitializeParallelDSM_hook = NULL;
383393
bbf_ParallelWorkerMain_hook = NULL;
394+
ParallelQueryMain_hook = prev_ParallelQueryMain_hook;
395+
ExecInitParallelPlan_hook = prev_ExecInitParallelPlan_hook;
396+
ExecCheckRTEPerms_hook = NULL;
384397
}
385398

386399
/*****************************************

‎contrib/babelfishpg_tsql/src/pltsql.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include "commands/trigger.h"
2727
#include "collation.h"
2828
#include "executor/spi.h"
29+
#include "libpq/libpq-be.h"
2930
#include "optimizer/planner.h"
3031
#include "utils/expandedrecord.h"
3132
#include "utils/plancache.h"
@@ -1772,6 +1773,8 @@ extern common_utility_plugin *common_utility_plugin_ptr;
17721773
#define IS_TDS_CLIENT() (*pltsql_protocol_plugin_ptr && \
17731774
(*pltsql_protocol_plugin_ptr)->is_tds_client)
17741775

1776+
#define IS_TDS_CONN() (MyProcPort && MyProcPort->is_tds_conn)
1777+
17751778
extern Oid procid_var;
17761779
extern uint64 rowcount_var;
17771780
extern List *columns_updated_list;

‎contrib/babelfishpg_tsql/src/session.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "pltsql.h"
1919
#include "guc.h"
2020
#include "storage/shm_toc.h"
21+
#include "bbf_parallel_query.h"
2122

2223
/* Core Session Properties */
2324

0 commit comments

Comments
 (0)