Fix rewriter to cope (more or less) with CTEs in the query being rewritten.

Since the original implementation of CTEs only allowed them in SELECT
queries, the rule rewriter did not expect to find any CTEs in statements
being rewritten by ON INSERT/UPDATE/DELETE rules.  We had dealt with this
to some extent but the code was still several bricks shy of a load, as
illustrated in bug #6051 from Jehan-Guillaume de Rorthais.

In particular, we have to be able to copy CTEs from the original query's
cteList into that of a rule action, in case the rule action references the
CTE (which it pretty much always will).  This also implies we were doing
things in the wrong order in RewriteQuery: we have to recursively rewrite
the CTE queries before expanding the main query, so that we have the
rewritten queries available to copy.

There are unpleasant limitations yet to resolve here, but at least we now
throw understandable FEATURE_NOT_SUPPORTED errors for them instead of just
failing with bizarre implementation-dependent errors.  In particular, we
can't handle propagating the same CTE into multiple post-rewrite queries
(because then the CTE would be evaluated multiple times), and we can't cope
with conflicts between CTE names in the original query and in the rule
actions.

3 files changed, 189 insertions, 61 deletions

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index bfc8fd7ee0..be9e7a4598 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -455,6 +455,44 @@ rewriteRuleAction(Query *parsetree,
 	}
 
 	/*
+	 * If the original query has any CTEs, copy them into the rule action.
+	 * But we don't need them for a utility action.
+	 */
+	if (parsetree->cteList != NIL && sub_action->commandType != CMD_UTILITY)
+	{
+		ListCell   *lc;
+
+		/*
+		 * Annoying implementation restriction: because CTEs are identified
+		 * by name within a cteList, we can't merge a CTE from the original
+		 * query if it has the same name as any CTE in the rule action.
+		 *
+		 * This could possibly be fixed by using some sort of internally
+		 * generated ID, instead of names, to link CTE RTEs to their CTEs.
+		 */
+		foreach(lc, parsetree->cteList)
+		{
+			CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc);
+			ListCell   *lc2;
+
+			foreach(lc2, sub_action->cteList)
+			{
+				CommonTableExpr *cte2 = (CommonTableExpr *) lfirst(lc2);
+
+				if (strcmp(cte->ctename, cte2->ctename) == 0)
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("WITH query name \"%s\" appears in both a rule action and the query being rewritten",
+									cte->ctename)));
+			}
+		}
+
+		/* OK, it's safe to combine the CTE lists */
+		sub_action->cteList = list_concat(sub_action->cteList,
+										  copyObject(parsetree->cteList));
+	}
+
+	/*
 	 * Event Qualification forces copying of parsetree and splitting into two
 	 * queries one w/rule_qual, one w/NOT rule_qual. Also add user query qual
 	 * onto rule action
@@ -1806,6 +1844,69 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 	ListCell   *lc1;
 
 	/*
+	 * First, recursively process any insert/update/delete statements in WITH
+	 * clauses.  (We have to do this first because the WITH clauses may get
+	 * copied into rule actions below.)
+	 */
+	foreach(lc1, parsetree->cteList)
+	{
+		CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc1);
+		Query	   *ctequery = (Query *) cte->ctequery;
+		List	   *newstuff;
+
+		Assert(IsA(ctequery, Query));
+
+		if (ctequery->commandType == CMD_SELECT)
+			continue;
+
+		newstuff = RewriteQuery(ctequery, rewrite_events);
+
+		/*
+		 * Currently we can only handle unconditional, single-statement DO
+		 * INSTEAD rules correctly; we have to get exactly one Query out of
+		 * the rewrite operation to stuff back into the CTE node.
+		 */
+		if (list_length(newstuff) == 1)
+		{
+			/* Push the single Query back into the CTE node */
+			ctequery = (Query *) linitial(newstuff);
+			Assert(IsA(ctequery, Query));
+			/* WITH queries should never be canSetTag */
+			Assert(!ctequery->canSetTag);
+			cte->ctequery = (Node *) ctequery;
+		}
+		else if (newstuff == NIL)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("DO INSTEAD NOTHING rules are not supported for data-modifying statements in WITH")));
+		}
+		else
+		{
+			ListCell   *lc2;
+
+			/* examine queries to determine which error message to issue */
+			foreach(lc2, newstuff)
+			{
+				Query	   *q = (Query *) lfirst(lc2);
+
+				if (q->querySource == QSRC_QUAL_INSTEAD_RULE)
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("conditional DO INSTEAD rules are not supported for data-modifying statements in WITH")));
+				if (q->querySource == QSRC_NON_INSTEAD_RULE)
+					ereport(ERROR,
+							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							 errmsg("DO ALSO rules are not supported for data-modifying statements in WITH")));
+			}
+
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("multi-statement DO INSTEAD rules are not supported for data-modifying statements in WITH")));
+		}
+	}
+
+	/*
 	 * If the statement is an insert, update, or delete, adjust its targetlist
 	 * as needed, and then fire INSERT/UPDATE/DELETE rules on it.
 	 *
@@ -1984,67 +2085,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 	}
 
 	/*
-	 * Recursively process any insert/update/delete statements in WITH clauses
-	 */
-	foreach(lc1, parsetree->cteList)
-	{
-		CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc1);
-		Query	   *ctequery = (Query *) cte->ctequery;
-		List	   *newstuff;
-
-		Assert(IsA(ctequery, Query));
-
-		if (ctequery->commandType == CMD_SELECT)
-			continue;
-
-		newstuff = RewriteQuery(ctequery, rewrite_events);
-
-		/*
-		 * Currently we can only handle unconditional, single-statement DO
-		 * INSTEAD rules correctly; we have to get exactly one Query out of
-		 * the rewrite operation to stuff back into the CTE node.
-		 */
-		if (list_length(newstuff) == 1)
-		{
-			/* Push the single Query back into the CTE node */
-			ctequery = (Query *) linitial(newstuff);
-			Assert(IsA(ctequery, Query));
-			/* WITH queries should never be canSetTag */
-			Assert(!ctequery->canSetTag);
-			cte->ctequery = (Node *) ctequery;
-		}
-		else if (newstuff == NIL)
-		{
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("DO INSTEAD NOTHING rules are not supported for data-modifying statements in WITH")));
-		}
-		else
-		{
-			ListCell   *lc2;
-
-			/* examine queries to determine which error message to issue */
-			foreach(lc2, newstuff)
-			{
-				Query	   *q = (Query *) lfirst(lc2);
-
-				if (q->querySource == QSRC_QUAL_INSTEAD_RULE)
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("conditional DO INSTEAD rules are not supported for data-modifying statements in WITH")));
-				if (q->querySource == QSRC_NON_INSTEAD_RULE)
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("DO ALSO rules are not supported for data-modifying statements in WITH")));
-			}
-
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("multi-statement DO INSTEAD rules are not supported for data-modifying statements in WITH")));
-		}
-	}
-
-	/*
 	 * For INSERTs, the original query is done first; for UPDATE/DELETE, it is
 	 * done last.  This is needed because update and delete rule actions might
 	 * not do anything if they are invoked after the update or delete is
@@ -2074,6 +2114,31 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 		}
 	}
 
+	/*
+	 * If the original query has a CTE list, and we generated more than one
+	 * non-utility result query, we have to fail because we'll have copied
+	 * the CTE list into each result query.  That would break the expectation
+	 * of single evaluation of CTEs.  This could possibly be fixed by
+	 * restructuring so that a CTE list can be shared across multiple Query
+	 * and PlannableStatement nodes.
+	 */
+	if (parsetree->cteList != NIL)
+	{
+		int		qcount = 0;
+
+		foreach(lc1, rewritten)
+		{
+			Query	   *q = (Query *) lfirst(lc1);
+
+			if (q->commandType != CMD_UTILITY)
+				qcount++;
+		}
+		if (qcount > 1)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("WITH cannot be used in a query that is rewritten by rules into multiple queries")));
+	}
+
 	return rewritten;
 }
 
diff --git a/src/test/regress/expected/with.out b/src/test/regress/expected/with.out
index b31d58f816..a1b089921d 100644
--- a/src/test/regress/expected/with.out
+++ b/src/test/regress/expected/with.out
@@ -1397,6 +1397,46 @@ SELECT * FROM y;
 (17 rows)
 
 DROP RULE y_rule ON y;
+-- check merging of outer CTE with CTE in a rule action
+CREATE TEMP TABLE bug6051 AS
+  select i from generate_series(1,3) as t(i);
+SELECT * FROM bug6051;
+ i 
+---
+ 1
+ 2
+ 3
+(3 rows)
+
+WITH t1 AS ( DELETE FROM bug6051 RETURNING * )
+INSERT INTO bug6051 SELECT * FROM t1;
+SELECT * FROM bug6051;
+ i 
+---
+ 1
+ 2
+ 3
+(3 rows)
+
+CREATE TEMP TABLE bug6051_2 (i int);
+CREATE RULE bug6051_ins AS ON INSERT TO bug6051 DO INSTEAD
+ INSERT INTO bug6051_2
+ SELECT NEW.i;
+WITH t1 AS ( DELETE FROM bug6051 RETURNING * )
+INSERT INTO bug6051 SELECT * FROM t1;
+SELECT * FROM bug6051;
+ i 
+---
+(0 rows)
+
+SELECT * FROM bug6051_2;
+ i 
+---
+ 1
+ 2
+ 3
+(3 rows)
+
 -- a truly recursive CTE in the same list
 WITH RECURSIVE t(a) AS (
 	SELECT 0
diff --git a/src/test/regress/sql/with.sql b/src/test/regress/sql/with.sql
index 4f6b517103..bc340e4543 100644
--- a/src/test/regress/sql/with.sql
+++ b/src/test/regress/sql/with.sql
@@ -610,6 +610,29 @@ SELECT * FROM y;
 
 DROP RULE y_rule ON y;
 
+-- check merging of outer CTE with CTE in a rule action
+CREATE TEMP TABLE bug6051 AS
+  select i from generate_series(1,3) as t(i);
+
+SELECT * FROM bug6051;
+
+WITH t1 AS ( DELETE FROM bug6051 RETURNING * )
+INSERT INTO bug6051 SELECT * FROM t1;
+
+SELECT * FROM bug6051;
+
+CREATE TEMP TABLE bug6051_2 (i int);
+
+CREATE RULE bug6051_ins AS ON INSERT TO bug6051 DO INSTEAD
+ INSERT INTO bug6051_2
+ SELECT NEW.i;
+
+WITH t1 AS ( DELETE FROM bug6051 RETURNING * )
+INSERT INTO bug6051 SELECT * FROM t1;
+
+SELECT * FROM bug6051;
+SELECT * FROM bug6051_2;
+
 -- a truly recursive CTE in the same list
 WITH RECURSIVE t(a) AS (
 	SELECT 0