make PrelExpressionForRelid() more memory-safe

Author: funbringer

src/include/relation_info.h

@@ -261,15 +261,11 @@ PrelExpressionColumnNames(const PartRelationInfo *prel)
 static inline Node *
 PrelExpressionForRelid(const PartRelationInfo *prel, Index rel_index)
 {
-	Node *expr;
-
 	/* TODO: implement some kind of cache */
+	Node *expr = copyObject(prel->expr);
+
 	if (rel_index != PART_EXPR_VARNO)
-	{
-		expr = copyObject(prel->expr);
 		ChangeVarNodes(expr, PART_EXPR_VARNO, rel_index, 0);
-	}
-	else expr = prel->expr;
 
 	return expr;
 }

src/include/runtimeappend.h

@@ -41,6 +41,9 @@ typedef struct
 	/* Refined clauses for partition pruning */
 	List			   *canon_custom_exprs;
 
+	/* Copy of partitioning expression (protect from invalidations) */
+	Node			   *prel_expr;
+
 	/* All available plans \ plan states */
 	HTAB			   *children_table;
 	HASHCTL				children_table_config;

src/nodes_common.c

@@ -635,10 +635,16 @@ create_append_scan_state_common(CustomScan *node,
 void
 begin_append_common(CustomScanState *node, EState *estate, int eflags)
 {
-	RuntimeAppendState *scan_state = (RuntimeAppendState *) node;
+	RuntimeAppendState	   *scan_state = (RuntimeAppendState *) node;
+	const PartRelationInfo *prel;
 
 	node->ss.ps.ps_TupFromTlist = false;
 
+	prel = get_pathman_relation_info(scan_state->relid);
+
+	/* Prepare expression according to set_set_customscan_references() */
+	scan_state->prel_expr = PrelExpressionForRelid(prel, INDEX_VAR);
+
 	/* Prepare custom expression according to set_set_customscan_references() */
 	scan_state->canon_custom_exprs =
 			canonicalize_custom_exprs(scan_state->custom_exprs);
@@ -709,18 +715,14 @@ rescan_append_common(CustomScanState *node)
 	WalkerContext			wcxt;
 	Oid					   *parts;
 	int						nparts;
-	Node				   *prel_expr;
 
 	prel = get_pathman_relation_info(scan_state->relid);
 	Assert(prel);
 
-	/* Prepare expression according to set_set_customscan_references() */
-	prel_expr = PrelExpressionForRelid(prel, INDEX_VAR);
-
 	/* First we select all available partitions... */
 	ranges = list_make1_irange_full(prel, IR_COMPLETE);
 
-	InitWalkerContext(&wcxt, prel_expr, prel, econtext);
+	InitWalkerContext(&wcxt, scan_state->prel_expr, prel, econtext);
 	foreach (lc, scan_state->canon_custom_exprs)
 	{
 		WrapperNode *wrap;

src/partition_filter.c

@@ -560,9 +560,6 @@ partition_filter_begin(CustomScanState *node, EState *estate, int eflags)
 		}
 		expr = PrelExpressionForRelid(prel, parent_varno);
 
-		/* HACK: protect expression from 'prel' invalidation */
-		expr = copyObject(expr);
-
 		/* Prepare state for expression execution */
 		old_mcxt = MemoryContextSwitchTo(estate->es_query_cxt);
 		state->expr_state = ExecInitExpr((Expr *) expr, NULL);