Support data-modifying commands (INSERT/UPDATE/DELETE) in WITH.

This patch implements data-modifying WITH queries according to the
semantics that the updates all happen with the same command counter value,
and in an unspecified order.  Therefore one WITH clause can't see the
effects of another, nor can the outer query see the effects other than
through the RETURNING values.  And attempts to do conflicting updates will
have unpredictable results.  We'll need to document all that.

This commit just fixes the code; documentation updates are waiting on
author.

Marko Tiikkaja and Hitoshi Harada

10 files changed, 36 insertions, 10 deletions

diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
index 0a0544bc38a..106c1eed2fa 100644
--- a/src/include/catalog/catversion.h
+++ b/src/include/catalog/catversion.h
@@ -53,6 +53,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	201102222
+#define CATALOG_VERSION_NO	201102251
 
 #endif
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 018b14ac843..2ed54d0a5cc 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -163,10 +163,10 @@ extern void ExecutorEnd(QueryDesc *queryDesc);
 extern void standard_ExecutorEnd(QueryDesc *queryDesc);
 extern void ExecutorRewind(QueryDesc *queryDesc);
 extern bool ExecCheckRTPerms(List *rangeTable, bool ereport_on_violation);
+extern void CheckValidResultRel(Relation resultRel, CmdType operation);
 extern void InitResultRelInfo(ResultRelInfo *resultRelInfo,
 				  Relation resultRelationDesc,
 				  Index resultRelationIndex,
-				  CmdType operation,
 				  int instrument_options);
 extern ResultRelInfo *ExecGetTriggerResultRel(EState *estate, Oid relid);
 extern bool ExecContextForcesOids(PlanState *planstate, bool *hasoids);
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 0a6b829de4f..d9aec4c26a3 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -346,7 +346,7 @@ typedef struct EState
 	/* If query can insert/delete tuples, the command ID to mark them with */
 	CommandId	es_output_cid;
 
-	/* Info about target table for insert/update/delete queries: */
+	/* Info about target table(s) for insert/update/delete queries: */
 	ResultRelInfo *es_result_relations; /* array of ResultRelInfos */
 	int			es_num_result_relations;		/* length of array */
 	ResultRelInfo *es_result_relation_info;		/* currently active array elt */
@@ -378,6 +378,8 @@ typedef struct EState
 
 	List	   *es_subplanstates;		/* List of PlanState for SubPlans */
 
+	List	   *es_auxmodifytables;	/* List of secondary ModifyTableStates */
+
 	/*
 	 * this ExprContext is for per-output-tuple operations, such as constraint
 	 * checks and index-value computations.  It will be reset for each output
@@ -1041,10 +1043,13 @@ typedef struct ResultState
 typedef struct ModifyTableState
 {
 	PlanState	ps;				/* its first field is NodeTag */
-	CmdType		operation;
+	CmdType		operation;		/* INSERT, UPDATE, or DELETE */
+	bool		canSetTag;		/* do we set the command tag/es_processed? */
+	bool		mt_done;		/* are we done? */
 	PlanState **mt_plans;		/* subplans (one per target rel) */
 	int			mt_nplans;		/* number of plans in the array */
 	int			mt_whichplan;	/* which one is being executed (0..n-1) */
+	ResultRelInfo *resultRelInfo;	/* per-subplan target relations */
 	List	  **mt_arowmarks;	/* per-subplan ExecAuxRowMark lists */
 	EPQState	mt_epqstate;	/* for evaluating EvalPlanQual rechecks */
 	bool		fireBSTriggers; /* do we need to fire stmt triggers? */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 536c03245e3..824403c69b3 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -118,6 +118,7 @@ typedef struct Query
 	bool		hasSubLinks;	/* has subquery SubLink */
 	bool		hasDistinctOn;	/* distinctClause is from DISTINCT ON */
 	bool		hasRecursive;	/* WITH RECURSIVE was specified */
+	bool		hasModifyingCTE;	/* has INSERT/UPDATE/DELETE in WITH */
 	bool		hasForUpdate;	/* FOR UPDATE or FOR SHARE was specified */
 
 	List	   *cteList;		/* WITH list (of CommonTableExpr's) */
@@ -884,7 +885,8 @@ typedef struct CommonTableExpr
 	NodeTag		type;
 	char	   *ctename;		/* query name (never qualified) */
 	List	   *aliascolnames;	/* optional list of column names */
-	Node	   *ctequery;		/* subquery (SelectStmt or Query) */
+	/* SelectStmt/InsertStmt/etc before parse analysis, Query afterwards: */
+	Node	   *ctequery;		/* the CTE's subquery */
 	int			location;		/* token location, or -1 if unknown */
 	/* These fields are set during parse analysis: */
 	bool		cterecursive;	/* is this CTE actually recursive? */
@@ -896,6 +898,14 @@ typedef struct CommonTableExpr
 	List	   *ctecolcollations; /* OID list of column collation OIDs */
 } CommonTableExpr;
 
+/* Convenience macro to get the output tlist of a CTE's query */
+#define GetCTETargetList(cte) \
+	(AssertMacro(IsA((cte)->ctequery, Query)), \
+	 ((Query *) (cte)->ctequery)->commandType == CMD_SELECT ? \
+	 ((Query *) (cte)->ctequery)->targetList : \
+	 ((Query *) (cte)->ctequery)->returningList)
+
+
 /*****************************************************************************
  *		Optimizable Statements
  *****************************************************************************/
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 90d61256e9c..efc79186f95 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -40,6 +40,8 @@ typedef struct PlannedStmt
 
 	bool		hasReturning;	/* is it insert|update|delete RETURNING? */
 
+	bool		hasModifyingCTE;	/* has insert|update|delete in WITH? */
+
 	bool		canSetTag;		/* do I set the command result tag? */
 
 	bool		transientPlan;	/* redo plan when TransactionXmin changes? */
@@ -167,7 +169,9 @@ typedef struct ModifyTable
 {
 	Plan		plan;
 	CmdType		operation;		/* INSERT, UPDATE, or DELETE */
+	bool		canSetTag;		/* do we set the command tag/es_processed? */
 	List	   *resultRelations;	/* integer list of RT indexes */
+	int			resultRelIndex;	/* index of first resultRel in plan's list */
 	List	   *plans;			/* plan(s) producing source data */
 	List	   *returningLists; /* per-target-table RETURNING tlists */
 	List	   *rowMarks;		/* PlanRowMarks (non-locking only) */
diff --git a/src/include/nodes/relation.h b/src/include/nodes/relation.h
index ab708351ed7..8bcc4006a1a 100644
--- a/src/include/nodes/relation.h
+++ b/src/include/nodes/relation.h
@@ -76,6 +76,8 @@ typedef struct PlannerGlobal
 
 	List	   *finalrowmarks;	/* "flat" list of PlanRowMarks */
 
+	List	   *resultRelations;	/* "flat" list of integer RT indexes */
+
 	List	   *relationOids;	/* OIDs of relations the plan depends on */
 
 	List	   *invalItems;		/* other dependencies, as PlanInvalItems */
@@ -154,8 +156,6 @@ typedef struct PlannerInfo
 	List	  **join_rel_level; /* lists of join-relation RelOptInfos */
 	int			join_cur_level; /* index of list being extended */
 
-	List	   *resultRelations;	/* integer list of RT indexes, or NIL */
-
 	List	   *init_plans;		/* init SubPlans for query */
 
 	List	   *cte_plan_ids;	/* per-CTE-item list of subplan IDs */
diff --git a/src/include/optimizer/planmain.h b/src/include/optimizer/planmain.h
index 9ddd5c183e3..7e03bc924ed 100644
--- a/src/include/optimizer/planmain.h
+++ b/src/include/optimizer/planmain.h
@@ -78,8 +78,8 @@ extern SetOp *make_setop(SetOpCmd cmd, SetOpStrategy strategy, Plan *lefttree,
 		   long numGroups, double outputRows);
 extern Result *make_result(PlannerInfo *root, List *tlist,
 			Node *resconstantqual, Plan *subplan);
-extern ModifyTable *make_modifytable(CmdType operation, List *resultRelations,
-				 List *subplans, List *returningLists,
+extern ModifyTable *make_modifytable(CmdType operation, bool canSetTag,
+				 List *resultRelations, List *subplans, List *returningLists,
 				 List *rowMarks, int epqParam);
 extern bool is_projection_capable_plan(Plan *plan);
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index a1511cbe64c..a68f7cf5087 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -103,6 +103,7 @@ struct ParseState
 	bool		p_hasAggs;
 	bool		p_hasWindowFuncs;
 	bool		p_hasSubLinks;
+	bool		p_hasModifyingCTE;
 	bool		p_is_insert;
 	bool		p_is_update;
 	bool		p_locked_from_parent;
diff --git a/src/include/parser/parse_relation.h b/src/include/parser/parse_relation.h
index 41f482c8df9..50ee4dacd28 100644
--- a/src/include/parser/parse_relation.h
+++ b/src/include/parser/parse_relation.h
@@ -74,7 +74,7 @@ extern RangeTblEntry *addRangeTableEntryForJoin(ParseState *pstate,
 extern RangeTblEntry *addRangeTableEntryForCTE(ParseState *pstate,
 						 CommonTableExpr *cte,
 						 Index levelsup,
-						 Alias *alias,
+						 RangeVar *rv,
 						 bool inFromCl);
 extern bool isLockedRefname(ParseState *pstate, const char *refname);
 extern void addRTEtoQuery(ParseState *pstate, RangeTblEntry *rte,
diff --git a/src/include/utils/portal.h b/src/include/utils/portal.h
index 639ccc45469..51e2bc99843 100644
--- a/src/include/utils/portal.h
+++ b/src/include/utils/portal.h
@@ -71,6 +71,11 @@
  * can't cope, and also because we don't want to risk failing to execute
  * all the auxiliary queries.)
  *
+ * PORTAL_ONE_MOD_WITH: the portal contains one single SELECT query, but
+ * it has data-modifying CTEs.  This is currently treated the same as the
+ * PORTAL_ONE_RETURNING case because of the possibility of needing to fire
+ * triggers.  It may act more like PORTAL_ONE_SELECT in future.
+ *
  * PORTAL_UTIL_SELECT: the portal contains a utility statement that returns
  * a SELECT-like result (for example, EXPLAIN or SHOW).  On first execution,
  * we run the statement and dump its results into the portal tuplestore;
@@ -83,6 +88,7 @@ typedef enum PortalStrategy
 {
 	PORTAL_ONE_SELECT,
 	PORTAL_ONE_RETURNING,
+	PORTAL_ONE_MOD_WITH,
 	PORTAL_UTIL_SELECT,
 	PORTAL_MULTI_QUERY
 } PortalStrategy;