1 files changed, 52 insertions, 5 deletions

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;
 }