diff --git a/pg_hint_plan.c b/pg_hint_plan.c index 0fc32cc6..1fc9471d 100644 --- a/pg_hint_plan.c +++ b/pg_hint_plan.c @@ -565,6 +565,7 @@ static planner_hook_type prev_planner = NULL; static join_search_hook_type prev_join_search = NULL; static set_rel_pathlist_hook_type prev_set_rel_pathlist = NULL; static ProcessUtility_hook_type prev_ProcessUtility_hook = NULL; +static ExecutorEnd_hook_type prev_ExecutorEnd = NULL; /* Hold reference to currently active hint */ static HintState *current_hint_state = NULL; @@ -618,6 +619,26 @@ PLpgSQL_plugin plugin_funcs = { NULL, }; +/* + * pg_hint_ExecutorEnd + * + * Force a hint to be retrieved when we are at the top of a PL recursion + * level. This can become necessary to handle hints in queries executed + * in the extended protocol, where the executor can be executed multiple + * times in a portal, but it could be possible to fail the hint retrieval. + */ +static void +pg_hint_ExecutorEnd(QueryDesc *queryDesc) +{ + if (plpgsql_recurse_level <= 0) + current_hint_retrieved = false; + + if (prev_ExecutorEnd) + prev_ExecutorEnd(queryDesc); + else + standard_ExecutorEnd(queryDesc); +} + /* * Module load callbacks */ @@ -698,6 +719,8 @@ _PG_init(void) set_rel_pathlist_hook = pg_hint_plan_set_rel_pathlist; prev_ProcessUtility_hook = ProcessUtility_hook; ProcessUtility_hook = pg_hint_plan_ProcessUtility; + prev_ExecutorEnd = ExecutorEnd_hook; + ExecutorEnd_hook = pg_hint_ExecutorEnd; /* setup PL/pgSQL plugin hook */ var_ptr = (PLpgSQL_plugin **) find_rendezvous_variable("PLpgSQL_plugin"); @@ -721,6 +744,7 @@ _PG_fini(void) join_search_hook = prev_join_search; set_rel_pathlist_hook = prev_set_rel_pathlist; ProcessUtility_hook = prev_ProcessUtility_hook; + ExecutorEnd_hook = prev_ExecutorEnd; /* uninstall PL/pgSQL plugin hook */ var_ptr = (PLpgSQL_plugin **) find_rendezvous_variable("PLpgSQL_plugin");