make PrelExpressionForRelid() more memory-safe

funbringer · funbringer · commit 0963406aeb82 · 2017-06-01T15:02:30.000+03:00
diff --git a/src/include/relation_info.h b/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;
 }
diff --git a/src/include/runtimeappend.h b/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;
diff --git a/src/nodes_common.c b/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;
diff --git a/src/partition_filter.c b/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);

Original file line number	Diff line number	Diff line change
`@@ -560,9 +560,6 @@ partition_filter_begin(CustomScanState node, EState estate, int eflags)`
`560`	`560`	`}`
`561`	`561`	`expr = PrelExpressionForRelid(prel, parent_varno);`
`562`	`562`
`563`		`- /* HACK: protect expression from 'prel' invalidation */`
`564`		`- expr = copyObject(expr);`
`565`		`-`
`566`	`563`	`/* Prepare state for expression execution */`
`567`	`564`	`old_mcxt = MemoryContextSwitchTo(estate->es_query_cxt);`
`568`	`565`	`state->expr_state = ExecInitExpr((Expr *) expr, NULL);`