3 files changed, 82 insertions, 17 deletions

diff --git a/src/backend/rewrite/rewriteDefine.c b/src/backend/rewrite/rewriteDefine.c
index 7ca90a54857..75653240c96 100644
--- a/src/backend/rewrite/rewriteDefine.c
+++ b/src/backend/rewrite/rewriteDefine.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteDefine.c,v 1.129 2008/08/25 22:42:34 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteDefine.c,v 1.130 2008/10/04 21:56:54 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -635,17 +635,23 @@ setRuleCheckAsUser_Query(Query *qry, Oid userid)
 			rte->checkAsUser = userid;
 	}
 
+	/* Recurse into subquery-in-WITH */
+	foreach(l, qry->cteList)
+	{
+		CommonTableExpr *cte = (CommonTableExpr *) lfirst(l);
+
+		setRuleCheckAsUser_Query((Query *) cte->ctequery, userid);
+	}
+
 	/* If there are sublinks, search for them and process their RTEs */
-	/* ignore subqueries in rtable because we already processed them */
 	if (qry->hasSubLinks)
 		query_tree_walker(qry, setRuleCheckAsUser_walker, (void *) &userid,
-						  QTW_IGNORE_RT_SUBQUERIES);
+						  QTW_IGNORE_RC_SUBQUERIES);
 }
 
 
 /*
  * Change the firing semantics of an existing rule.
- *
  */
 void
 EnableDisableRule(Relation rel, const char *rulename,
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 305b91202ad..f07be31e3ac 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -7,7 +7,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteHandler.c,v 1.180 2008/09/24 16:52:46 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteHandler.c,v 1.181 2008/10/04 21:56:54 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -215,13 +215,21 @@ AcquireRewriteLocks(Query *parsetree)
 		}
 	}
 
+	/* Recurse into subqueries in WITH */
+	foreach(l, parsetree->cteList)
+	{
+		CommonTableExpr *cte = (CommonTableExpr *) lfirst(l);
+
+		AcquireRewriteLocks((Query *) cte->ctequery);
+	}
+
 	/*
 	 * Recurse into sublink subqueries, too.  But we already did the ones in
-	 * the rtable.
+	 * the rtable and cteList.
 	 */
 	if (parsetree->hasSubLinks)
 		query_tree_walker(parsetree, acquireLocksOnSubLinks, NULL,
-						  QTW_IGNORE_RT_SUBQUERIES);
+						  QTW_IGNORE_RC_SUBQUERIES);
 }
 
 /*
@@ -1228,6 +1236,35 @@ markQueryForLocking(Query *qry, Node *jtnode, bool forUpdate, bool noWait)
 			markQueryForLocking(rte->subquery, (Node *) rte->subquery->jointree,
 								forUpdate, noWait);
 		}
+		else if (rte->rtekind == RTE_CTE)
+		{
+			/*
+			 * We allow FOR UPDATE/SHARE of a WITH query to be propagated into
+			 * the WITH, but it doesn't seem very sane to allow this for a
+			 * reference to an outer-level WITH (compare
+			 * transformLockingClause).  Which simplifies life here.
+			 */
+			CommonTableExpr *cte = NULL;
+			ListCell   *lc;
+
+			if (rte->ctelevelsup > 0 || rte->self_reference)
+				ereport(ERROR,
+						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						 errmsg("SELECT FOR UPDATE/SHARE cannot be applied to an outer-level WITH query")));
+			foreach(lc, qry->cteList)
+			{
+				cte = (CommonTableExpr *) lfirst(lc);
+				if (strcmp(cte->ctename, rte->ctename) == 0)
+					break;
+			}
+			if (lc == NULL)				/* shouldn't happen */
+				elog(ERROR, "could not find CTE \"%s\"", rte->ctename);
+			/* should be analyzed by now */
+			Assert(IsA(cte->ctequery, Query));
+			markQueryForLocking((Query *) cte->ctequery,
+								(Node *) ((Query *) cte->ctequery)->jointree,
+								forUpdate, noWait);
+		}
 	}
 	else if (IsA(jtnode, FromExpr))
 	{
@@ -1295,6 +1332,7 @@ static Query *
 fireRIRrules(Query *parsetree, List *activeRIRs)
 {
 	int			rt_index;
+	ListCell   *lc;
 
 	/*
 	 * don't try to convert this into a foreach loop, because rtable list can
@@ -1407,13 +1445,22 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		heap_close(rel, NoLock);
 	}
 
+	/* Recurse into subqueries in WITH */
+	foreach(lc, parsetree->cteList)
+	{
+		CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc);
+
+		cte->ctequery = (Node *)
+			fireRIRrules((Query *) cte->ctequery, activeRIRs);
+	}
+
 	/*
 	 * Recurse into sublink subqueries, too.  But we already did the ones in
-	 * the rtable.
+	 * the rtable and cteList.
 	 */
 	if (parsetree->hasSubLinks)
 		query_tree_walker(parsetree, fireRIRonSubLink, (void *) activeRIRs,
-						  QTW_IGNORE_RT_SUBQUERIES);
+						  QTW_IGNORE_RC_SUBQUERIES);
 
 	return parsetree;
 }
diff --git a/src/backend/rewrite/rewriteManip.c b/src/backend/rewrite/rewriteManip.c
index c8652ff5f2d..b5d82dd0e7f 100644
--- a/src/backend/rewrite/rewriteManip.c
+++ b/src/backend/rewrite/rewriteManip.c
@@ -7,7 +7,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteManip.c,v 1.113 2008/09/01 20:42:45 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/rewrite/rewriteManip.c,v 1.114 2008/10/04 21:56:54 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -183,13 +183,13 @@ bool
 checkExprHasSubLink(Node *node)
 {
 	/*
-	 * If a Query is passed, examine it --- but we need not recurse into
-	 * sub-Queries.
+	 * If a Query is passed, examine it --- but we should not recurse into
+	 * sub-Queries that are in its rangetable or CTE list.
 	 */
 	return query_or_expression_tree_walker(node,
 										   checkExprHasSubLink_walker,
 										   NULL,
-										   QTW_IGNORE_RT_SUBQUERIES);
+										   QTW_IGNORE_RC_SUBQUERIES);
 }
 
 static bool
@@ -543,7 +543,7 @@ adjust_relid_set(Relids relids, int oldrelid, int newrelid)
  * that sublink are not affected, only outer references to vars that belong
  * to the expression's original query level or parents thereof.
  *
- * Aggref nodes are adjusted similarly.
+ * Likewise for other nodes containing levelsup fields, such as Aggref.
  *
  * NOTE: although this has the form of a walker, we cheat and modify the
  * Var nodes in-place.	The given expression tree should have been copied
@@ -585,6 +585,17 @@ IncrementVarSublevelsUp_walker(Node *node,
 			agg->agglevelsup += context->delta_sublevels_up;
 		/* fall through to recurse into argument */
 	}
+	if (IsA(node, RangeTblEntry))
+	{
+		RangeTblEntry *rte = (RangeTblEntry *) node;
+
+		if (rte->rtekind == RTE_CTE)
+		{
+			if (rte->ctelevelsup >= context->min_sublevels_up)
+				rte->ctelevelsup += context->delta_sublevels_up;
+		}
+		return false;			/* allow range_table_walker to continue */
+	}
 	if (IsA(node, Query))
 	{
 		/* Recurse into subselects */
@@ -593,7 +604,8 @@ IncrementVarSublevelsUp_walker(Node *node,
 		context->min_sublevels_up++;
 		result = query_tree_walker((Query *) node,
 								   IncrementVarSublevelsUp_walker,
-								   (void *) context, 0);
+								   (void *) context,
+								   QTW_EXAMINE_RTES);
 		context->min_sublevels_up--;
 		return result;
 	}
@@ -617,7 +629,7 @@ IncrementVarSublevelsUp(Node *node, int delta_sublevels_up,
 	query_or_expression_tree_walker(node,
 									IncrementVarSublevelsUp_walker,
 									(void *) &context,
-									0);
+									QTW_EXAMINE_RTES);
 }
 
 /*
@@ -636,7 +648,7 @@ IncrementVarSublevelsUp_rtable(List *rtable, int delta_sublevels_up,
 	range_table_walker(rtable,
 					   IncrementVarSublevelsUp_walker,
 					   (void *) &context,
-					   0);
+					   QTW_EXAMINE_RTES);
 }