1 files changed, 55 insertions, 41 deletions

diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 0865e3ebef2..a717a0deead 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -1768,7 +1768,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	Oid			my_lastoid = InvalidOid;
 	SPITupleTable *my_tuptable = NULL;
 	int			res = 0;
-	bool		have_active_snap = ActiveSnapshotSet();
+	bool		pushed_active_snap = false;
 	ErrorContextCallback spierrcontext;
 	CachedPlan *cplan = NULL;
 	ListCell   *lc1;
@@ -1781,6 +1781,40 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 	spierrcontext.previous = error_context_stack;
 	error_context_stack = &spierrcontext;
 
+	/*
+	 * We support four distinct snapshot management behaviors:
+	 *
+	 * snapshot != InvalidSnapshot, read_only = true: use exactly the given
+	 * snapshot.
+	 *
+	 * snapshot != InvalidSnapshot, read_only = false: use the given
+	 * snapshot, modified by advancing its command ID before each querytree.
+	 *
+	 * snapshot == InvalidSnapshot, read_only = true: use the entry-time
+	 * ActiveSnapshot, if any (if there isn't one, we run with no snapshot).
+	 *
+	 * snapshot == InvalidSnapshot, read_only = false: take a full new
+	 * snapshot for each user command, and advance its command ID before each
+	 * querytree within the command.
+	 *
+	 * In the first two cases, we can just push the snap onto the stack
+	 * once for the whole plan list.
+	 */
+	if (snapshot != InvalidSnapshot)
+	{
+		if (read_only)
+		{
+			PushActiveSnapshot(snapshot);
+			pushed_active_snap = true;
+		}
+		else
+		{
+			/* Make sure we have a private copy of the snapshot to modify */
+			PushCopiedSnapshot(snapshot);
+			pushed_active_snap = true;
+		}
+	}
+
 	foreach(lc1, plan->plancache_list)
 	{
 		CachedPlanSource *plansource = (CachedPlanSource *) lfirst(lc1);
@@ -1802,12 +1836,23 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 			stmt_list = plansource->plan->stmt_list;
 		}
 
+		/*
+		 * In the default non-read-only case, get a new snapshot, replacing
+		 * any that we pushed in a previous cycle.
+		 */
+		if (snapshot == InvalidSnapshot && !read_only)
+		{
+			if (pushed_active_snap)
+				PopActiveSnapshot();
+			PushActiveSnapshot(GetTransactionSnapshot());
+			pushed_active_snap = true;
+		}
+
 		foreach(lc2, stmt_list)
 		{
 			Node	   *stmt = (Node *) lfirst(lc2);
 			bool		canSetTag;
 			DestReceiver *dest;
-			bool		pushed_active_snap = false;
 
 			_SPI_current->processed = 0;
 			_SPI_current->lastoid = InvalidOid;
@@ -1848,48 +1893,16 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 
 			/*
 			 * If not read-only mode, advance the command counter before each
-			 * command.
+			 * command and update the snapshot.
 			 */
 			if (!read_only)
+			{
 				CommandCounterIncrement();
+				UpdateActiveSnapshotCommandId();
+			}
 
 			dest = CreateDestReceiver(canSetTag ? DestSPI : DestNone);
 
-			if (snapshot == InvalidSnapshot)
-			{
-				/*
-				 * Default read_only behavior is to use the entry-time
-				 * ActiveSnapshot, if any; if read-write, grab a full new
-				 * snap.
-				 */
-				if (read_only)
-				{
-					if (have_active_snap)
-					{
-						PushActiveSnapshot(GetActiveSnapshot());
-						pushed_active_snap = true;
-					}
-				}
-				else
-				{
-					PushActiveSnapshot(GetTransactionSnapshot());
-					pushed_active_snap = true;
-				}
-			}
-			else
-			{
-				/*
-				 * We interpret read_only with a specified snapshot to be
-				 * exactly that snapshot, but read-write means use the snap
-				 * with advancing of command ID.
-				 */
-				if (read_only)
-					PushActiveSnapshot(snapshot);
-				else
-					PushUpdatedSnapshot(snapshot);
-				pushed_active_snap = true;
-			}
-
 			if (IsA(stmt, PlannedStmt) &&
 				((PlannedStmt *) stmt)->utilityStmt == NULL)
 			{
@@ -1925,9 +1938,6 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 				res = SPI_OK_UTILITY;
 			}
 
-			if (pushed_active_snap)
-				PopActiveSnapshot();
-
 			/*
 			 * The last canSetTag query sets the status values returned to the
 			 * caller.	Be careful to free any tuptables not returned, to
@@ -1970,6 +1980,10 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 
 fail:
 
+	/* Pop the snapshot off the stack if we pushed one */
+	if (pushed_active_snap)
+		PopActiveSnapshot();
+
 	/* We no longer need the cached plan refcount, if any */
 	if (cplan)
 		ReleaseCachedPlan(cplan, true);