[PGPRO-7630] Post-processing for nodes added in plan tree by pathman

New nodes added in pathman planner hook had no correct plan_node_id,
which could cause problems later for statistics collector.
Added fixes for 'custom_scan_tlist' to let EXPLAIN (VERBOSE) work.
Also changed queryId type on uint64.
Added hook for compatibility with pgpro_stats.
Fixed tree walkers for ModifyTable.

Tags: pg_pathman

Author: Svetlana Derevyanko

Diff for: src/hooks.c

@@ -21,6 +21,7 @@
 #include "hooks.h"
 #include "init.h"
 #include "partition_filter.h"
+#include "partition_overseer.h"
 #include "partition_router.h"
 #include "pathman_workers.h"
 #include "planner_tree_modification.h"
@@ -74,6 +75,7 @@ planner_hook_type				pathman_planner_hook_next				= NULL;
 post_parse_analyze_hook_type	pathman_post_parse_analyze_hook_next	= NULL;
 shmem_startup_hook_type			pathman_shmem_startup_hook_next			= NULL;
 ProcessUtility_hook_type		pathman_process_utility_hook_next		= NULL;
+ExecutorStart_hook_type			pathman_executor_start_hook_prev		= NULL;
 
 
 /* Take care of joins */
@@ -673,6 +675,23 @@ execute_for_plantree(PlannedStmt *planned_stmt,
 	planned_stmt->subplans = subplans;
 }
 
+/*
+ * Truncated version of set_plan_refs.
+ * Pathman can add nodes to already completed and post-processed plan tree.
+ * reset_plan_node_ids fixes some presentation values for updated plan tree
+ * to avoid problems in further processing.
+ */
+static Plan *
+reset_plan_node_ids(Plan *plan, void *lastPlanNodeId)
+{
+	if (plan == NULL)
+		return NULL;
+
+	plan->plan_node_id = (*(int *) lastPlanNodeId)++;
+
+	return plan;
+}
+
 /*
  * Planner hook. It disables inheritance for tables that have been partitioned
  * by pathman to prevent standart PostgreSQL partitioning mechanism from
@@ -688,7 +707,7 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 #endif
 {
 	PlannedStmt	   *result;
-	uint32			query_id = parse->queryId;
+	uint64			query_id = parse->queryId;
 
 	/* Save the result in case it changes */
 	bool			pathman_ready = IsPathmanReady();
@@ -720,6 +739,9 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 
 		if (pathman_ready)
 		{
+			int		lastPlanNodeId = 0;
+			ListCell *l;
+
 			/* Add PartitionFilter node for INSERT queries */
 			execute_for_plantree(result, add_partition_filters);
 
@@ -729,6 +751,13 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 			/* Decrement planner() calls count */
 			decr_planner_calls_count();
 
+			/* remake parsed tree presentation fixes due to possible adding nodes */
+			result->planTree = plan_tree_visitor(result->planTree, reset_plan_node_ids, &lastPlanNodeId);
+			foreach(l, result->subplans)
+			{
+				lfirst(l) = plan_tree_visitor((Plan *) lfirst(l), reset_plan_node_ids, &lastPlanNodeId);
+			}
+
 			/* HACK: restore queryId set by pg_stat_statements */
 			result->queryId = query_id;
 		}
@@ -1125,3 +1154,53 @@ pathman_process_utility_hook(Node *first_arg,
 								dest, completionTag);
 #endif
 }
+
+/*
+ * Planstate tree nodes could have been copied.
+ * It breaks references on correspoding
+ * ModifyTable node from PartitionRouter nodes.
+ */
+static void
+fix_mt_refs(PlanState *state, void *context)
+{
+	ModifyTableState		   *mt_state = (ModifyTableState *) state;
+	PartitionRouterState	   *pr_state;
+#if PG_VERSION_NUM < 140000
+	int							i;
+#endif
+
+	if (!IsA(state, ModifyTableState))
+		return;
+#if PG_VERSION_NUM >= 140000
+	{
+		CustomScanState *pf_state = (CustomScanState *) outerPlanState(mt_state);
+#else
+	for (i = 0; i < mt_state->mt_nplans; i++)
+	{
+		CustomScanState *pf_state = (CustomScanState *) mt_state->mt_plans[i];
+#endif
+		if (IsPartitionFilterState(pf_state))
+		{
+			pr_state = linitial(pf_state->custom_ps);
+			if (IsPartitionRouterState(pr_state))
+			{
+				pr_state->mt_state = mt_state;
+			}
+		}
+	}
+}
+
+void
+pathman_executor_start_hook(QueryDesc *queryDesc, int eflags)
+{
+	if (pathman_executor_start_hook_prev)
+		pathman_executor_start_hook_prev(queryDesc, eflags);
+	else
+		standard_ExecutorStart(queryDesc, eflags);
+
+	/*
+	 * HACK for compatibility with pgpro_stats.
+	 * Fix possibly broken planstate tree.
+	 */
+	state_tree_visitor(queryDesc->planstate, fix_mt_refs, NULL);
+}

Diff for: src/include/hooks.h

@@ -28,6 +28,7 @@ extern post_parse_analyze_hook_type		pathman_post_parse_analyze_hook_next;
 extern shmem_startup_hook_type			pathman_shmem_startup_hook_next;
 extern ProcessUtility_hook_type			pathman_process_utility_hook_next;
 extern ExecutorRun_hook_type			pathman_executor_run_hook_next;
+extern ExecutorStart_hook_type			pathman_executor_start_hook_prev;
 
 
 void pathman_join_pathlist_hook(PlannerInfo *root,
@@ -115,4 +116,6 @@ void pathman_executor_hook(QueryDesc *queryDesc,
 						   ExecutorRun_CountArgType count);
 #endif
 
+void pathman_executor_start_hook(QueryDesc *queryDescc,
+								 int eflags);
 #endif /* PATHMAN_HOOKS_H */

Diff for: src/include/partition_filter.h

@@ -183,6 +183,7 @@ void destroy_tuple_map(TupleConversionMap *tuple_map);
 
 List * pfilter_build_tlist(Plan *subplan);
 
+void pfilter_tlist_fix_resjunk(CustomScan *subplan);
 
 /* Find suitable partition using 'value' */
 Oid * find_partitions_for_value(Datum value, Oid value_type,

Diff for: src/include/partition_router.h

@@ -78,7 +78,7 @@ void partition_router_explain(CustomScanState *node,
 							  List *ancestors,
 							  ExplainState *es);
 
-Plan *make_partition_router(Plan *subplan, int epq_param);
+Plan *make_partition_router(Plan *subplan, int epq_param, Index parent_rti);
 Node *partition_router_create_scan_state(CustomScan *node);
 TupleTableSlot *partition_router_exec(CustomScanState *node);
 

Diff for: src/partition_filter.c

@@ -817,6 +817,7 @@ make_partition_filter(Plan *subplan,
 	/* Prepare 'custom_scan_tlist' for EXPLAIN (VERBOSE) */
 	cscan->custom_scan_tlist = copyObject(cscan->scan.plan.targetlist);
 	ChangeVarNodes((Node *) cscan->custom_scan_tlist, INDEX_VAR, parent_rti, 0);
+	pfilter_tlist_fix_resjunk(cscan);
 
 	/* Pack partitioned table's Oid and conflict_action */
 	cscan->custom_private = list_make4(makeInteger(parent_relid),
@@ -1114,6 +1115,34 @@ pfilter_build_tlist(Plan *subplan)
 	return result_tlist;
 }
 
+/*
+ * resjunk Vars had its varattnos being set on nonexisting relation columns.
+ * For future processing service attributes should be indicated correctly.
+ */
+void
+pfilter_tlist_fix_resjunk(CustomScan *css)
+{
+	ListCell *lc;
+
+	foreach(lc, css->custom_scan_tlist)
+	{
+		TargetEntry *tle = (TargetEntry *) lfirst(lc);
+
+		if (!IsA(tle->expr, Const))
+		{
+			Var *var = (Var *) tle->expr;
+
+			if (tle->resjunk)
+			{
+				/* To make Var recognizable as service attribute. */
+				var->varattno = -1;
+			}
+		}
+	}
+
+	return;
+}
+
 /*
  * ----------------------------------------------
  *  Additional init steps for ResultPartsStorage

Diff for: src/partition_router.c

@@ -115,7 +115,7 @@ init_partition_router_static_data(void)
 }
 
 Plan *
-make_partition_router(Plan *subplan, int epq_param)
+make_partition_router(Plan *subplan, int epq_param, Index parent_rti)
 {
 	CustomScan *cscan = makeNode(CustomScan);
 
@@ -136,8 +136,10 @@ make_partition_router(Plan *subplan, int epq_param)
 	/* Build an appropriate target list */
 	cscan->scan.plan.targetlist = pfilter_build_tlist(subplan);
 
-	/* FIXME: should we use the same tlist? */
-	cscan->custom_scan_tlist = subplan->targetlist;
+	/* Fix 'custom_scan_tlist' for EXPLAIN (VERBOSE) */
+	cscan->custom_scan_tlist = copyObject(cscan->scan.plan.targetlist);
+	ChangeVarNodes((Node *) cscan->custom_scan_tlist, INDEX_VAR, parent_rti, 0);
+	pfilter_tlist_fix_resjunk(cscan);
 
 	return &cscan->scan.plan;
 }

Diff for: src/pg_pathman.c

@@ -357,6 +357,8 @@ _PG_init(void)
 	planner_hook							= pathman_planner_hook;
 	pathman_process_utility_hook_next		= ProcessUtility_hook;
 	ProcessUtility_hook						= pathman_process_utility_hook;
+	pathman_executor_start_hook_prev		= ExecutorStart_hook;
+	ExecutorStart_hook						= pathman_executor_start_hook;
 
 	/* Initialize static data for all subsystems */
 	init_main_pathman_toggles();

Diff for: src/planner_tree_modification.c

@@ -122,8 +122,8 @@ static void handle_modification_query(Query *parse, transform_query_cxt *context
 static Plan *partition_filter_visitor(Plan *plan, void *context);
 static Plan *partition_router_visitor(Plan *plan, void *context);
 
-static void state_visit_subplans(List *plans, void (*visitor) (), void *context);
-static void state_visit_members(PlanState **planstates, int nplans, void (*visitor) (), void *context);
+static void state_visit_subplans(List *plans, void (*visitor) (PlanState *plan, void *context), void *context);
+static void state_visit_members(PlanState **planstates, int nplans, void (*visitor) (PlanState *plan, void *context), void *context);
 
 static Oid find_deepest_partition(Oid relid, Index rti, Expr *quals);
 static Node *eval_extern_params_mutator(Node *node, ParamListInfo params);
@@ -137,13 +137,13 @@ static bool modifytable_contains_fdw(List *rtable, ModifyTable *node);
  * id in order to recognize them properly.
  */
 #define QUERY_ID_INITIAL 0
-static uint32 latest_query_id = QUERY_ID_INITIAL;
+static uint64 latest_query_id = QUERY_ID_INITIAL;
 
 
 void
 assign_query_id(Query *query)
 {
-	uint32	prev_id = latest_query_id++;
+	uint64	prev_id = latest_query_id++;
 
 	if (prev_id > latest_query_id)
 		elog(WARNING, "assign_query_id(): queryId overflow");
@@ -187,14 +187,12 @@ plan_tree_visitor(Plan *plan,
 				plan_tree_visitor((Plan *) lfirst(l), visitor, context);
 			break;
 
+#if PG_VERSION_NUM < 140000 /* reworked in commit 86dc90056dfd */
 		case T_ModifyTable:
-#if PG_VERSION_NUM >= 140000 /* reworked in commit 86dc90056dfd */
-			plan_tree_visitor(outerPlan(plan), visitor, context);
-#else
 			foreach (l, ((ModifyTable *) plan)->plans)
 				plan_tree_visitor((Plan *) lfirst(l), visitor, context);
-#endif
 			break;
+#endif
 
 		case T_Append:
 			foreach (l, ((Append *) plan)->appendplans)
@@ -254,15 +252,13 @@ state_tree_visitor(PlanState *state,
 				state_tree_visitor((PlanState *) lfirst(lc), visitor, context);
 			break;
 
+#if PG_VERSION_NUM < 140000 /* reworked in commit 86dc90056dfd */
 		case T_ModifyTable:
-#if PG_VERSION_NUM >= 140000 /* reworked in commit 86dc90056dfd */
-			visitor(outerPlanState(state), context);
-#else
 			state_visit_members(((ModifyTableState *) state)->mt_plans,
 								((ModifyTableState *) state)->mt_nplans,
 								visitor, context);
-#endif
 			break;
+#endif
 
 		case T_Append:
 			state_visit_members(((AppendState *) state)->appendplans,
@@ -307,15 +303,15 @@ state_tree_visitor(PlanState *state,
  */
 static void
 state_visit_subplans(List *plans,
-					 void (*visitor) (),
+					 void (*visitor) (PlanState *plan, void *context),
 					 void *context)
 {
 	ListCell *lc;
 
 	foreach (lc, plans)
 	{
 		SubPlanState *sps = lfirst_node(SubPlanState, lc);
-		visitor(sps->planstate, context);
+		state_tree_visitor(sps->planstate, visitor, context);
 	}
 }
 
@@ -325,12 +321,12 @@ state_visit_subplans(List *plans,
  */
 static void
 state_visit_members(PlanState **planstates, int nplans,
-					void (*visitor) (), void *context)
+					void (*visitor) (PlanState *plan, void *context), void *context)
 {
 	int i;
 
 	for (i = 0; i < nplans; i++)
-		visitor(planstates[i], context);
+		state_tree_visitor(planstates[i], visitor, context);
 }
 
 
@@ -939,10 +935,12 @@ partition_router_visitor(Plan *plan, void *context)
 
 #if PG_VERSION_NUM >= 140000 /* for changes 86dc90056dfd */
 			prouter = make_partition_router(subplan,
-											modify_table->epqParam);
+											modify_table->epqParam,
+											modify_table->nominalRelation);
 #else
 			prouter = make_partition_router((Plan *) lfirst(lc1),
-											modify_table->epqParam);
+											modify_table->epqParam,
+											modify_table->nominalRelation);
 #endif
 
 			pfilter = make_partition_filter((Plan *) prouter, relid,