Fix oversight in initial implementation of PORTAL_ONE_RETURNING mode: we

cannot assume that there's exactly one Query in the Portal, as we can for
ONE_SELECT mode, because non-SELECT queries might have extra queries added
during rule rewrites.  Fix things up so that we'll use ONE_RETURNING mode
when a Portal contains one primary (canSetTag) query and that query has
a RETURNING list.  This appears to be a second showstopper reason for running
the Portal to completion before we start to hand anything back --- we want
to be sure that the rule-added queries get run too.

5 files changed, 137 insertions, 63 deletions

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index b96426856cd..82ad85a410a 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -10,7 +10,7 @@
  * Copyright (c) 2002-2006, PostgreSQL Global Development Group
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/commands/prepare.c,v 1.60 2006/08/12 02:52:04 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/commands/prepare.c,v 1.61 2006/08/14 22:57:15 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -448,7 +448,7 @@ FetchPreparedStatementResultDesc(PreparedStatement *stmt)
 			return ExecCleanTypeFromTL(query->targetList, false);
 
 		case PORTAL_ONE_RETURNING:
-			query = (Query *) linitial(stmt->query_list);
+			query = PortalListGetPrimaryQuery(stmt->query_list);
 			return ExecCleanTypeFromTL(query->returningList, false);
 
 		case PORTAL_UTIL_SELECT:
@@ -505,7 +505,7 @@ FetchPreparedStatementTargetList(PreparedStatement *stmt)
 	if (strategy == PORTAL_ONE_SELECT)
 		return ((Query *) linitial(stmt->query_list))->targetList;
 	if (strategy == PORTAL_ONE_RETURNING)
-		return ((Query *) linitial(stmt->query_list))->returningList;
+		return (PortalListGetPrimaryQuery(stmt->query_list))->returningList;
 	if (strategy == PORTAL_UTIL_SELECT)
 	{
 		Node	   *utilityStmt;
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 165ff962fcd..183969f1b85 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/executor/spi.c,v 1.156 2006/08/14 13:40:18 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/executor/spi.c,v 1.157 2006/08/14 22:57:15 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -818,31 +818,44 @@ SPI_cursor_open(const char *name, void *plan,
 				bool read_only)
 {
 	_SPI_plan  *spiplan = (_SPI_plan *) plan;
-	List	   *qtlist = spiplan->qtlist;
-	List	   *ptlist = spiplan->ptlist;
-	Query	   *queryTree;
-	Plan	   *planTree;
+	List	   *qtlist;
+	List	   *ptlist;
 	ParamListInfo paramLI;
 	Snapshot	snapshot;
 	MemoryContext oldcontext;
 	Portal		portal;
 	int			k;
 
-	/* Ensure that the plan contains only one query */
-	if (list_length(ptlist) != 1 || list_length(qtlist) != 1)
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_CURSOR_DEFINITION),
-				 errmsg("cannot open multi-query plan as cursor")));
-	queryTree = (Query *) linitial((List *) linitial(qtlist));
-	planTree = (Plan *) linitial(ptlist);
-
-	/* Must be a query that returns tuples */
-	if (!QueryReturnsTuples(queryTree))
+	/*
+	 * Check that the plan is something the Portal code will special-case
+	 * as returning one tupleset.
+	 */
+	if (!SPI_is_cursor_plan(spiplan))
+	{
+		/* try to give a good error message */
+		Query	   *queryTree;
+
+		if (list_length(spiplan->qtlist) != 1)
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_CURSOR_DEFINITION),
+					 errmsg("cannot open multi-query plan as cursor")));
+		queryTree = PortalListGetPrimaryQuery((List *) linitial(spiplan->qtlist));
+		if (queryTree == NULL)
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_CURSOR_DEFINITION),
+					 errmsg("cannot open empty query as cursor")));
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_CURSOR_DEFINITION),
 				 /* translator: %s is name of a SQL command, eg INSERT */
 				 errmsg("cannot open %s query as cursor",
 						CreateQueryTag(queryTree))));
+	}
+
+	Assert(list_length(spiplan->qtlist) == 1);
+	qtlist = (List *) linitial(spiplan->qtlist);
+	ptlist = spiplan->ptlist;
+	if (list_length(qtlist) != list_length(ptlist))
+		elog(ERROR, "corrupted SPI plan lists");
 
 	/* Reset SPI result (note we deliberately don't touch lastoid) */
 	SPI_processed = 0;
@@ -862,10 +875,10 @@ SPI_cursor_open(const char *name, void *plan,
 		portal = CreatePortal(name, false, false);
 	}
 
-	/* Switch to portal's memory and copy the parsetree and plan to there */
+	/* Switch to portal's memory and copy the parsetrees and plans to there */
 	oldcontext = MemoryContextSwitchTo(PortalGetHeapMemory(portal));
-	queryTree = copyObject(queryTree);
-	planTree = copyObject(planTree);
+	qtlist = copyObject(qtlist);
+	ptlist = copyObject(ptlist);
 
 	/* If the plan has parameters, set them up */
 	if (spiplan->nargs > 0)
@@ -907,9 +920,9 @@ SPI_cursor_open(const char *name, void *plan,
 	PortalDefineQuery(portal,
 					  NULL,		/* no statement name */
 					  spiplan->query,
-					  CreateQueryTag(queryTree),
-					  list_make1(queryTree),
-					  list_make1(planTree),
+					  CreateQueryTag(PortalListGetPrimaryQuery(qtlist)),
+					  qtlist,
+					  ptlist,
 					  PortalGetHeapMemory(portal));
 
 	MemoryContextSwitchTo(oldcontext);
@@ -918,7 +931,8 @@ SPI_cursor_open(const char *name, void *plan,
 	 * Set up options for portal.
 	 */
 	portal->cursorOptions &= ~(CURSOR_OPT_SCROLL | CURSOR_OPT_NO_SCROLL);
-	if (planTree == NULL || ExecSupportsBackwardScan(planTree))
+	if (list_length(ptlist) == 1 &&
+		ExecSupportsBackwardScan((Plan *) linitial(ptlist)))
 		portal->cursorOptions |= CURSOR_OPT_SCROLL;
 	else
 		portal->cursorOptions |= CURSOR_OPT_NO_SCROLL;
@@ -940,16 +954,7 @@ SPI_cursor_open(const char *name, void *plan,
 	 */
 	PortalStart(portal, paramLI, snapshot);
 
-	/*
-	 * If this test fails then we're out of sync with pquery.c about
-	 * which queries can return tuples...
-	 */
-	if (portal->strategy == PORTAL_MULTI_QUERY)
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_CURSOR_DEFINITION),
-				 /* translator: %s is name of a SQL command, eg INSERT */
-				 errmsg("cannot open %s query as cursor",
-						CreateQueryTag(queryTree))));
+	Assert(portal->strategy != PORTAL_MULTI_QUERY);
 
 	/* Return the created portal */
 	return portal;
@@ -1050,7 +1055,6 @@ bool
 SPI_is_cursor_plan(void *plan)
 {
 	_SPI_plan  *spiplan = (_SPI_plan *) plan;
-	List	   *qtlist;
 
 	if (spiplan == NULL)
 	{
@@ -1058,13 +1062,20 @@ SPI_is_cursor_plan(void *plan)
 		return false;
 	}
 
-	qtlist = spiplan->qtlist;
-	if (list_length(spiplan->ptlist) == 1 && list_length(qtlist) == 1)
-	{
-		Query	   *queryTree = (Query *) linitial((List *) linitial(qtlist));
+	if (list_length(spiplan->qtlist) != 1)
+		return false;			/* not exactly 1 pre-rewrite command */
 
-		if (QueryReturnsTuples(queryTree))
+	switch (ChoosePortalStrategy((List *) linitial(spiplan->qtlist)))
+	{
+		case PORTAL_ONE_SELECT:
+		case PORTAL_ONE_RETURNING:
+		case PORTAL_UTIL_SELECT:
+			/* OK */
 			return true;
+
+		case PORTAL_MULTI_QUERY:
+			/* will not return tuples */
+			break;
 	}
 	return false;
 }
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 2d12b0e7c4f..3ede40570ab 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/tcop/pquery.c,v 1.106 2006/08/12 02:52:05 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/backend/tcop/pquery.c,v 1.107 2006/08/14 22:57:15 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -213,30 +213,59 @@ ProcessQuery(Query *parsetree,
 PortalStrategy
 ChoosePortalStrategy(List *parseTrees)
 {
-	PortalStrategy strategy;
-
-	strategy = PORTAL_MULTI_QUERY;		/* default assumption */
+	int			nSetTag;
+	ListCell   *lc;
 
+	/*
+	 * PORTAL_ONE_SELECT and PORTAL_UTIL_SELECT need only consider the
+	 * single-Query-struct case, since there are no rewrite rules that
+	 * can add auxiliary queries to a SELECT or a utility command.
+	 */
 	if (list_length(parseTrees) == 1)
 	{
 		Query	   *query = (Query *) linitial(parseTrees);
 
+		Assert(IsA(query, Query));
 		if (query->canSetTag)
 		{
 			if (query->commandType == CMD_SELECT &&
 				query->into == NULL)
-				strategy = PORTAL_ONE_SELECT;
-			else if (query->returningList != NIL)
-				strategy = PORTAL_ONE_RETURNING;
-			else if (query->commandType == CMD_UTILITY &&
-					 query->utilityStmt != NULL)
+				return PORTAL_ONE_SELECT;
+			if (query->commandType == CMD_UTILITY &&
+				query->utilityStmt != NULL)
 			{
 				if (UtilityReturnsTuples(query->utilityStmt))
-					strategy = PORTAL_UTIL_SELECT;
+					return PORTAL_UTIL_SELECT;
+				/* it can't be ONE_RETURNING, so give up */
+				return PORTAL_MULTI_QUERY;
 			}
 		}
 	}
-	return strategy;
+
+	/*
+	 * PORTAL_ONE_RETURNING has to allow auxiliary queries added by rewrite.
+	 * Choose PORTAL_ONE_RETURNING if there is exactly one canSetTag query
+	 * and it has a RETURNING list.
+	 */
+	nSetTag = 0;
+	foreach(lc, parseTrees)
+	{
+		Query	   *query = (Query *) lfirst(lc);
+
+		Assert(IsA(query, Query));
+		if (query->canSetTag)
+		{
+			if (++nSetTag > 1)
+				return PORTAL_MULTI_QUERY;	/* no need to look further */
+			if (query->returningList == NIL)
+				return PORTAL_MULTI_QUERY;	/* no need to look further */
+		}
+	}
+	if (nSetTag == 1)
+		return PORTAL_ONE_RETURNING;
+
+	/* Else, it's the general case... */
+	return PORTAL_MULTI_QUERY;
 }
 
 /*
@@ -255,7 +284,7 @@ FetchPortalTargetList(Portal portal)
 	if (portal->strategy == PORTAL_ONE_SELECT)
 		return ((Query *) linitial(portal->parseTrees))->targetList;
 	if (portal->strategy == PORTAL_ONE_RETURNING)
-		return ((Query *) linitial(portal->parseTrees))->returningList;
+		return (PortalGetPrimaryQuery(portal))->returningList;
 	if (portal->strategy == PORTAL_UTIL_SELECT)
 	{
 		Node	   *utilityStmt;
@@ -422,7 +451,7 @@ PortalStart(Portal portal, ParamListInfo params, Snapshot snapshot)
 				 * the portal.  We do need to set up the result tupdesc.
 				 */
 				portal->tupDesc =
-					ExecCleanTypeFromTL(((Query *) linitial(portal->parseTrees))->returningList, false);
+					ExecCleanTypeFromTL((PortalGetPrimaryQuery(portal))->returningList, false);
 
 				/*
 				 * Reset cursor position data to "start of query"
@@ -894,10 +923,11 @@ FillPortalStore(Portal portal)
 	{
 		case PORTAL_ONE_RETURNING:
 			/*
-			 * We run the query just as if it were in a MULTI portal,
-			 * but send the output to the tuplestore.
+			 * Run the portal to completion just as for the default MULTI_QUERY
+			 * case, but send the primary query's output to the tuplestore.
+			 * Auxiliary query outputs are discarded.
 			 */
-			PortalRunMulti(portal, treceiver, treceiver, completionTag);
+			PortalRunMulti(portal, treceiver, None_Receiver, completionTag);
 			/* Override default completion tag with actual command result */
 			portal->commandTag = pstrdup(completionTag);
 			break;
diff --git a/src/backend/utils/mmgr/portalmem.c b/src/backend/utils/mmgr/portalmem.c
index da4aee4d162..85a6711d1cd 100644
--- a/src/backend/utils/mmgr/portalmem.c
+++ b/src/backend/utils/mmgr/portalmem.c
@@ -12,7 +12,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/utils/mmgr/portalmem.c,v 1.91 2006/08/08 01:23:15 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/backend/utils/mmgr/portalmem.c,v 1.92 2006/08/14 22:57:15 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -147,6 +147,34 @@ GetPortalByName(const char *name)
 }
 
 /*
+ * PortalListGetPrimaryQuery
+ *		Get the "primary" Query within a portal, ie, the one marked canSetTag.
+ *
+ * Returns NULL if no such Query.  If multiple Query structs within the
+ * portal are marked canSetTag, returns the first one.  Neither of these
+ * cases should occur in present usages of this function.
+ *
+ * Note: the reason this is just handed a List is so that prepared statements
+ * can share the code.  For use with a portal, use PortalGetPrimaryQuery
+ * rather than calling this directly.
+ */
+Query *
+PortalListGetPrimaryQuery(List *parseTrees)
+{
+	ListCell   *lc;
+
+	foreach(lc, parseTrees)
+	{
+		Query	   *query = (Query *) lfirst(lc);
+
+		Assert(IsA(query, Query));
+		if (query->canSetTag)
+			return query;
+	}
+	return NULL;
+}
+
+/*
  * CreatePortal
  *		Returns a new portal given a name.
  *
diff --git a/src/include/utils/portal.h b/src/include/utils/portal.h
index 3f308f5f582..92cc55189ed 100644
--- a/src/include/utils/portal.h
+++ b/src/include/utils/portal.h
@@ -39,7 +39,7 @@
  * Portions Copyright (c) 1996-2006, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/include/utils/portal.h,v 1.65 2006/08/12 02:52:06 tgl Exp $
+ * $PostgreSQL: pgsql/src/include/utils/portal.h,v 1.66 2006/08/14 22:57:15 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -63,10 +63,13 @@
  * tuplestore for access after transaction completion).
  *
  * PORTAL_ONE_RETURNING: the portal contains a single INSERT/UPDATE/DELETE
- * query with a RETURNING clause.  On first execution, we run the statement
- * and dump its results into the portal tuplestore; the results are then
- * returned to the client as demanded.  (We can't support suspension of
- * the query partway through, because the AFTER TRIGGER code can't cope.)
+ * query with a RETURNING clause (plus possibly auxiliary queries added by
+ * rule rewriting).  On first execution, we run the portal to completion
+ * and dump the primary query's results into the portal tuplestore; the
+ * results are then returned to the client as demanded.  (We can't support
+ * suspension of the query partway through, because the AFTER TRIGGER code
+ * can't cope, and also because we don't want to risk failing to execute
+ * all the auxiliary queries.)
  *
  * PORTAL_UTIL_SELECT: the portal contains a utility statement that returns
  * a SELECT-like result (for example, EXPLAIN or SHOW).  On first execution,
@@ -187,6 +190,7 @@ typedef struct PortalData
  */
 #define PortalGetQueryDesc(portal)	((portal)->queryDesc)
 #define PortalGetHeapMemory(portal) ((portal)->heap)
+#define PortalGetPrimaryQuery(portal) PortalListGetPrimaryQuery((portal)->parseTrees)
 
 
 /* Prototypes for functions in utils/mmgr/portalmem.c */
@@ -215,6 +219,7 @@ extern void PortalDefineQuery(Portal portal,
 				  List *parseTrees,
 				  List *planTrees,
 				  MemoryContext queryContext);
+extern Query *PortalListGetPrimaryQuery(List *parseTrees);
 extern void PortalCreateHoldStore(Portal portal);
 
 #endif   /* PORTAL_H */