From eb9d28bd28ad8741cacb6edca83d30b6da3d6589 Mon Sep 17 00:00:00 2001 From: Tom Lane Date: Fri, 14 Mar 2025 16:05:25 -0400 Subject: [PATCH v8 2/4] Use custom plan machinery for SQL function --- src/backend/executor/functions.c | 219 +++++++++++++++++----- src/test/regress/expected/rowsecurity.out | 51 +++++ src/test/regress/expected/rules.out | 35 ++++ src/test/regress/sql/rowsecurity.sql | 41 ++++ src/test/regress/sql/rules.sql | 24 +++ 5 files changed, 328 insertions(+), 42 deletions(-) diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c index 6aa8e9c4d8a..ae0425c050a 100644 --- a/src/backend/executor/functions.c +++ b/src/backend/executor/functions.c @@ -33,6 +33,7 @@ #include "utils/datum.h" #include "utils/lsyscache.h" #include "utils/memutils.h" +#include "utils/plancache.h" #include "utils/snapmgr.h" #include "utils/syscache.h" @@ -112,6 +113,12 @@ typedef struct JunkFilter *junkFilter; /* will be NULL if function returns VOID */ + /* Cached plans support */ + List *plansource_list; /* list of plansource */ + List *cplan_list; /* list of cached plans */ + int planning_stmt_number; /* the number of statement we are + * currently planning */ + /* * func_state is a List of execution_state records, each of which is the * first for its original parsetree, with any additional records chained @@ -122,6 +129,8 @@ typedef struct MemoryContext fcontext; /* memory context holding this struct and all * subsidiary data */ + MemoryContext planning_context; /* memory context which is used for + * planning */ LocalTransactionId lxid; /* lxid in which cache was made */ SubTransactionId subxid; /* subxid in which cache was made */ @@ -138,10 +147,9 @@ static Node *sql_fn_make_param(SQLFunctionParseInfoPtr pinfo, int paramno, int location); static Node *sql_fn_resolve_param_name(SQLFunctionParseInfoPtr pinfo, const char *paramname, int location); -static List *init_execution_state(List *queryTree_list, - SQLFunctionCachePtr fcache, +static List *init_execution_state(SQLFunctionCachePtr fcache, bool lazyEvalOK); -static void init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK); +static void init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool *lazyEvalOK); static void postquel_start(execution_state *es, SQLFunctionCachePtr fcache); static bool postquel_getnext(execution_state *es, SQLFunctionCachePtr fcache); static void postquel_end(execution_state *es); @@ -461,45 +469,52 @@ sql_fn_resolve_param_name(SQLFunctionParseInfoPtr pinfo, * querytrees. The sublist structure denotes the original query boundaries. */ static List * -init_execution_state(List *queryTree_list, - SQLFunctionCachePtr fcache, +init_execution_state(SQLFunctionCachePtr fcache, bool lazyEvalOK) { List *eslist = NIL; + List *cplan_list = NIL; execution_state *lasttages = NULL; ListCell *lc1; + MemoryContext oldcontext; + + /* + * Invalidate func_state prior to resetting - otherwise error callback can + * access it + */ + fcache->func_state = NIL; + MemoryContextReset(fcache->planning_context); + + oldcontext = MemoryContextSwitchTo(fcache->planning_context); - foreach(lc1, queryTree_list) + foreach(lc1, fcache->plansource_list) { - List *qtlist = lfirst_node(List, lc1); + CachedPlanSource *plansource = (CachedPlanSource *) lfirst(lc1); execution_state *firstes = NULL; execution_state *preves = NULL; ListCell *lc2; + CachedPlan *cplan; + + /* Save statement number for error reporting */ + fcache->planning_stmt_number = foreach_current_index(lc1) + 1; + + /* + * Get plan for the query. If paramLI is set, we can get custom plan + */ + cplan = GetCachedPlan(plansource, + fcache->paramLI, + plansource->is_saved ? CurrentResourceOwner : NULL, + NULL); - foreach(lc2, qtlist) + /* Record cplan in plan list to be released on replanning */ + cplan_list = lappend(cplan_list, cplan); + + /* For each planned statement create execution state */ + foreach(lc2, cplan->stmt_list) { - Query *queryTree = lfirst_node(Query, lc2); - PlannedStmt *stmt; + PlannedStmt *stmt = lfirst_node(PlannedStmt, lc2); execution_state *newes; - /* Plan the query if needed */ - if (queryTree->commandType == CMD_UTILITY) - { - /* Utility commands require no planning. */ - stmt = makeNode(PlannedStmt); - stmt->commandType = CMD_UTILITY; - stmt->canSetTag = queryTree->canSetTag; - stmt->utilityStmt = queryTree->utilityStmt; - stmt->stmt_location = queryTree->stmt_location; - stmt->stmt_len = queryTree->stmt_len; - stmt->queryId = queryTree->queryId; - } - else - stmt = pg_plan_query(queryTree, - fcache->src, - CURSOR_OPT_PARALLEL_OK, - NULL); - /* * Precheck all commands for validity in a function. This should * generally match the restrictions spi.c applies. @@ -541,7 +556,7 @@ init_execution_state(List *queryTree_list, newes->stmt = stmt; newes->qd = NULL; - if (queryTree->canSetTag) + if (stmt->canSetTag) lasttages = newes; preves = newes; @@ -573,6 +588,11 @@ init_execution_state(List *queryTree_list, fcache->lazyEval = lasttages->lazyEval = true; } + /* We've finished planning, reset planning statement number */ + fcache->planning_stmt_number = 0; + fcache->cplan_list = cplan_list; + + MemoryContextSwitchTo(oldcontext); return eslist; } @@ -580,7 +600,7 @@ init_execution_state(List *queryTree_list, * Initialize the SQLFunctionCache for a SQL function */ static void -init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) +init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool *lazyEvalOK) { FmgrInfo *finfo = fcinfo->flinfo; Oid foid = finfo->fn_oid; @@ -596,6 +616,7 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) ListCell *lc; Datum tmp; bool isNull; + List *plansource_list; /* * Create memory context that holds all the SQLFunctionCache data. It @@ -614,6 +635,10 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) */ fcache = (SQLFunctionCachePtr) palloc0(sizeof(SQLFunctionCache)); fcache->fcontext = fcontext; + /* Create separate context for planning */ + fcache->planning_context = AllocSetContextCreate(fcache->fcontext, + "SQL language functions planning context", + ALLOCSET_SMALL_SIZES); finfo->fn_extra = fcache; /* @@ -680,6 +705,7 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) * plancache.c. */ queryTree_list = NIL; + plansource_list = NIL; if (!isNull) { Node *n; @@ -695,8 +721,13 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) { Query *parsetree = lfirst_node(Query, lc); List *queryTree_sublist; + CachedPlanSource *plansource; AcquireRewriteLocks(parsetree, true, false); + + plansource = CreateCachedPlanForQuery(parsetree, fcache->src, CreateCommandTag((Node *) parsetree)); + plansource_list = lappend(plansource_list, plansource); + queryTree_sublist = pg_rewrite_query(parsetree); queryTree_list = lappend(queryTree_list, queryTree_sublist); } @@ -711,6 +742,10 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) { RawStmt *parsetree = lfirst_node(RawStmt, lc); List *queryTree_sublist; + CachedPlanSource *plansource; + + plansource = CreateCachedPlan(parsetree, fcache->src, CreateCommandTag(parsetree->stmt)); + plansource_list = lappend(plansource_list, plansource); queryTree_sublist = pg_analyze_and_rewrite_withcb(parsetree, fcache->src, @@ -751,6 +786,33 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) false, &resulttlist); + /* + * Queries could be rewritten by check_sql_fn_retval(). Now when they have + * their final form, we can complete plan cache entry creation. + */ + if (plansource_list != NIL) + { + ListCell *qlc; + ListCell *plc; + + forboth(qlc, queryTree_list, plc, plansource_list) + { + List *queryTree_sublist = lfirst(qlc); + CachedPlanSource *plansource = lfirst(plc); + + /* Finish filling in the CachedPlanSource */ + CompleteCachedPlan(plansource, + queryTree_sublist, + NULL, + NULL, + 0, + (ParserSetupHook) sql_fn_parser_setup, + fcache->pinfo, + CURSOR_OPT_PARALLEL_OK | CURSOR_OPT_NO_SCROLL, + false); + } + } + /* * Construct a JunkFilter we can use to coerce the returned rowtype to the * desired form, unless the result type is VOID, in which case there's @@ -792,13 +854,10 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK) * materialize mode, but to add more smarts in init_execution_state * about this, we'd probably need a three-way flag instead of bool. */ - lazyEvalOK = true; + *lazyEvalOK = true; } - /* Finally, plan the queries */ - fcache->func_state = init_execution_state(queryTree_list, - fcache, - lazyEvalOK); + fcache->plansource_list = plansource_list; /* Mark fcache with time of creation to show it's valid */ fcache->lxid = MyProc->vxid.lxid; @@ -971,7 +1030,12 @@ postquel_sub_params(SQLFunctionCachePtr fcache, prm->value = MakeExpandedObjectReadOnly(fcinfo->args[i].value, prm->isnull, get_typlen(argtypes[i])); - prm->pflags = 0; + + /* + * PARAM_FLAG_CONST is necessary to build efficient custom plan. + */ + prm->pflags = PARAM_FLAG_CONST; + prm->ptype = argtypes[i]; } } @@ -1024,6 +1088,33 @@ postquel_get_single_result(TupleTableSlot *slot, return value; } +/* + * Release plans. This function is called prior to planning + * statements with new parameters. When custom plans are generated + * for each function call in a statement, they can consume too much memory, so + * release them. Generic plans will survive it as plansource holds + * reference to a generic plan. + */ +static void +release_plans(List *cplans) +{ + ListCell *lc; + + /* + * We support separate plan list, so that we visit each plan here only + * once + */ + foreach(lc, cplans) + { + CachedPlan *cplan = lfirst(lc); + + ReleaseCachedPlan(cplan, cplan->is_saved ? CurrentResourceOwner : NULL); + } + + /* Cleanup the list itself */ + list_free(cplans); +} + /* * fmgr_sql: function call manager for SQL functions */ @@ -1042,6 +1133,7 @@ fmgr_sql(PG_FUNCTION_ARGS) Datum result; List *eslist; ListCell *eslc; + bool build_cached_plans = false; /* * Setup error traceback support for ereport() @@ -1097,7 +1189,11 @@ fmgr_sql(PG_FUNCTION_ARGS) if (fcache == NULL) { - init_sql_fcache(fcinfo, PG_GET_COLLATION(), lazyEvalOK); + /* + * init_sql_fcache() can set lazyEvalOK in additional cases when it + * determines that materialize won't work. + */ + init_sql_fcache(fcinfo, PG_GET_COLLATION(), &lazyEvalOK); fcache = (SQLFunctionCachePtr) fcinfo->flinfo->fn_extra; } @@ -1131,12 +1227,37 @@ fmgr_sql(PG_FUNCTION_ARGS) break; } + /* + * We skip actual planning for initial run, so in this case we have to + * build cached plans now. + */ + if (fcache->plansource_list != NIL && eslist == NIL) + build_cached_plans = true; + /* * Convert params to appropriate format if starting a fresh execution. (If * continuing execution, we can re-use prior params.) */ - if (is_first && es && es->status == F_EXEC_START) + if ((is_first && es && es->status == F_EXEC_START) || build_cached_plans) + { postquel_sub_params(fcache, fcinfo); + if (fcache->plansource_list) + { + /* replan the queries */ + fcache->func_state = init_execution_state(fcache, + lazyEvalOK); + /* restore execution state and eslist-related variables */ + eslist = fcache->func_state; + /* find the first non-NULL execution state */ + foreach(eslc, eslist) + { + es = (execution_state *) lfirst(eslc); + + if (es) + break; + } + } + } /* * Build tuplestore to hold results, if we don't have one already. Note @@ -1391,6 +1512,10 @@ fmgr_sql(PG_FUNCTION_ARGS) es = es->next; } } + + /* Release plans when functions stops executing */ + release_plans(fcache->cplan_list); + fcache->cplan_list = NULL; } error_context_stack = sqlerrcontext.previous; @@ -1430,13 +1555,19 @@ sql_exec_error_callback(void *arg) } /* - * Try to determine where in the function we failed. If there is a query - * with non-null QueryDesc, finger it. (We check this rather than looking - * for F_EXEC_RUN state, so that errors during ExecutorStart or + * Try to determine where in the function we failed. If failure happens + * while building plans, look at planning_stmt_number. Else if there is a + * query with non-null QueryDesc, finger it. (We check this rather than + * looking for F_EXEC_RUN state, so that errors during ExecutorStart or * ExecutorEnd are blamed on the appropriate query; see postquel_start and * postquel_end.) */ - if (fcache->func_state) + if (fcache->planning_stmt_number) + { + errcontext("SQL function \"%s\" statement %d", + fcache->fname, fcache->planning_stmt_number); + } + else if (fcache->func_state) { execution_state *es; int query_num; @@ -1522,6 +1653,10 @@ ShutdownSQLFunction(Datum arg) tuplestore_end(fcache->tstore); fcache->tstore = NULL; + /* Release plans when functions stops executing */ + release_plans(fcache->cplan_list); + fcache->cplan_list = NULL; + /* execUtils will deregister the callback... */ fcache->shutdown_reg = false; } diff --git a/src/test/regress/expected/rowsecurity.out b/src/test/regress/expected/rowsecurity.out index 87929191d06..438eaf69928 100644 --- a/src/test/regress/expected/rowsecurity.out +++ b/src/test/regress/expected/rowsecurity.out @@ -4695,6 +4695,57 @@ RESET ROLE; DROP FUNCTION rls_f(); DROP VIEW rls_v; DROP TABLE rls_t; +-- RLS changes invalidate cached function plans +create table rls_t (c text); +create table test_t (c text); +insert into rls_t values ('a'), ('b'), ('c'), ('d'); +insert into test_t values ('a'), ('b'); +alter table rls_t enable row level security; +grant select on rls_t to regress_rls_alice; +grant select on test_t to regress_rls_alice; +create policy p1 on rls_t for select to regress_rls_alice using (c = current_setting('rls_test.blah')); +-- Function changes row_security setting and so invalidates plan +create or replace function rls_f(text) + RETURNS text + LANGUAGE sql +BEGIN ATOMIC + select set_config('rls_test.blah', $1, true) || set_config('row_security', 'false', true) || string_agg(c, ',' order by c) from rls_t; +END; +-- Table owner bypasses RLS +select rls_f(c) from test_t order by rls_f; + rls_f +------------- + aoffa,b,c,d + boffa,b,c,d +(2 rows) + +set role regress_rls_alice; +-- For casual user changes in row_security setting lead +-- to error during query rewrite +select rls_f(c) from test_t order by rls_f; +ERROR: query would be affected by row-level security policy for table "rls_t" +CONTEXT: SQL function "rls_f" statement 1 +reset role; +set plan_cache_mode to force_generic_plan; +-- Table owner bypasses RLS, but cached plan invalidates +select rls_f(c) from test_t order by rls_f; + rls_f +------------- + aoffa,b,c,d + boffa,b,c,d +(2 rows) + +-- For casual user changes in row_security setting lead +-- to plan invalidation and error during query rewrite +set role regress_rls_alice; +select rls_f(c) from test_t order by rls_f; +ERROR: query would be affected by row-level security policy for table "rls_t" +CONTEXT: SQL function "rls_f" statement 1 +reset role; +reset plan_cache_mode; +reset rls_test.blah; +drop function rls_f; +drop table rls_t, test_t; -- -- Clean up objects -- diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out index 62f69ac20b2..b9fe71f391d 100644 --- a/src/test/regress/expected/rules.out +++ b/src/test/regress/expected/rules.out @@ -3878,3 +3878,38 @@ DROP TABLE ruletest_t3; DROP TABLE ruletest_t2; DROP TABLE ruletest_t1; DROP USER regress_rule_user1; +-- Test that SQL functions correctly handle DO NOTHING rule +CREATE TABLE some_data (i int, data text); +CREATE TABLE some_data_values (i int, data text); +CREATE FUNCTION insert_data(i int, data text) +RETURNS INT +AS $$ +INSERT INTO some_data VALUES ($1, $2); +SELECT 1; +$$ LANGUAGE SQL; +INSERT INTO some_data_values SELECT i , 'data'|| i FROM generate_series(1, 10) i; +CREATE RULE some_data_noinsert AS ON INSERT TO some_data DO INSTEAD NOTHING; +SELECT insert_data(i, data) FROM some_data_values; + insert_data +------------- + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 +(10 rows) + +SELECT * FROM some_data ORDER BY i; + i | data +---+------ +(0 rows) + +DROP RULE some_data_noinsert ON some_data; +DROP TABLE some_data_values; +DROP TABLE some_data; +DROP FUNCTION insert_data(int, text); diff --git a/src/test/regress/sql/rowsecurity.sql b/src/test/regress/sql/rowsecurity.sql index f61dbbf9581..9fe8f4b059c 100644 --- a/src/test/regress/sql/rowsecurity.sql +++ b/src/test/regress/sql/rowsecurity.sql @@ -2307,6 +2307,47 @@ DROP FUNCTION rls_f(); DROP VIEW rls_v; DROP TABLE rls_t; +-- RLS changes invalidate cached function plans +create table rls_t (c text); +create table test_t (c text); + +insert into rls_t values ('a'), ('b'), ('c'), ('d'); +insert into test_t values ('a'), ('b'); +alter table rls_t enable row level security; +grant select on rls_t to regress_rls_alice; +grant select on test_t to regress_rls_alice; +create policy p1 on rls_t for select to regress_rls_alice using (c = current_setting('rls_test.blah')); + +-- Function changes row_security setting and so invalidates plan +create or replace function rls_f(text) + RETURNS text + LANGUAGE sql +BEGIN ATOMIC + select set_config('rls_test.blah', $1, true) || set_config('row_security', 'false', true) || string_agg(c, ',' order by c) from rls_t; +END; + +-- Table owner bypasses RLS +select rls_f(c) from test_t order by rls_f; +set role regress_rls_alice; +-- For casual user changes in row_security setting lead +-- to error during query rewrite +select rls_f(c) from test_t order by rls_f; +reset role; + +set plan_cache_mode to force_generic_plan; +-- Table owner bypasses RLS, but cached plan invalidates +select rls_f(c) from test_t order by rls_f; +-- For casual user changes in row_security setting lead +-- to plan invalidation and error during query rewrite +set role regress_rls_alice; +select rls_f(c) from test_t order by rls_f; +reset role; +reset plan_cache_mode; +reset rls_test.blah; + +drop function rls_f; +drop table rls_t, test_t; + -- -- Clean up objects -- diff --git a/src/test/regress/sql/rules.sql b/src/test/regress/sql/rules.sql index fdd3ff1d161..505449452ee 100644 --- a/src/test/regress/sql/rules.sql +++ b/src/test/regress/sql/rules.sql @@ -1432,3 +1432,27 @@ DROP TABLE ruletest_t2; DROP TABLE ruletest_t1; DROP USER regress_rule_user1; + +-- Test that SQL functions correctly handle DO NOTHING rule +CREATE TABLE some_data (i int, data text); +CREATE TABLE some_data_values (i int, data text); + +CREATE FUNCTION insert_data(i int, data text) +RETURNS INT +AS $$ +INSERT INTO some_data VALUES ($1, $2); +SELECT 1; +$$ LANGUAGE SQL; + +INSERT INTO some_data_values SELECT i , 'data'|| i FROM generate_series(1, 10) i; + +CREATE RULE some_data_noinsert AS ON INSERT TO some_data DO INSTEAD NOTHING; + +SELECT insert_data(i, data) FROM some_data_values; + +SELECT * FROM some_data ORDER BY i; + +DROP RULE some_data_noinsert ON some_data; +DROP TABLE some_data_values; +DROP TABLE some_data; +DROP FUNCTION insert_data(int, text); -- 2.43.5