|
| 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 | + |
0 commit comments