From 8d40409d8c8e80e53ff4a06eb9ac99c5f30fb427 Mon Sep 17 00:00:00 2001
From: Kyotaro Horiguchi <horikyota.ntt@gmail.com>
Date: Tue, 29 Mar 2022 17:53:39 +0900
Subject: [PATCH v2] Appropriately set snapshot on cursor fetch

FETCH runs with the registered snapshot at DECLARE.  So if the result
contained TOASTed values, the same snapshot is required.  This patch
make PortalRunSelect to pass the DECLARE'd query's snapshot to outer
portal so that RunFromStore can use the snapshot.

Reported-by: Erik Rijkers <er@xs4all.nl>
---
 src/backend/tcop/dest.c               |  8 ++--
 src/backend/tcop/pquery.c             | 65 ++++++++++++++++++++++++++-
 src/include/tcop/dest.h               |  4 +-
 src/test/regress/expected/portals.out |  9 ++++
 src/test/regress/sql/portals.sql      | 10 +++++
 5 files changed, 90 insertions(+), 6 deletions(-)

diff --git a/src/backend/tcop/dest.c b/src/backend/tcop/dest.c
index c952cbea8b..4a8a67323e 100644
--- a/src/backend/tcop/dest.c
+++ b/src/backend/tcop/dest.c
@@ -69,22 +69,22 @@ donothingCleanup(DestReceiver *self)
  */
 static const DestReceiver donothingDR = {
 	donothingReceive, donothingStartup, donothingCleanup, donothingCleanup,
-	DestNone
+	NULL, DestNone
 };
 
 static const DestReceiver debugtupDR = {
 	debugtup, debugStartup, donothingCleanup, donothingCleanup,
-	DestDebug
+	NULL, DestDebug
 };
 
 static const DestReceiver printsimpleDR = {
 	printsimple, printsimple_startup, donothingCleanup, donothingCleanup,
-	DestRemoteSimple
+	NULL, DestRemoteSimple
 };
 
 static const DestReceiver spi_printtupDR = {
 	spi_printtup, spi_dest_startup, donothingCleanup, donothingCleanup,
-	DestSPI
+	NULL, DestSPI
 };
 
 /*
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 5aa5a350f3..001ff30db3 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -27,7 +27,7 @@
 #include "utils/memutils.h"
 #include "utils/snapmgr.h"
 
-
+extern const char *resownname(ResourceOwner o);
 /*
  * ActivePortal is the currently executing Portal (the most closely nested,
  * if there are several).
@@ -925,6 +925,26 @@ PortalRunSelect(Portal portal,
 						portal->run_once);
 			nprocessed = queryDesc->estate->es_processed;
 			PopActiveSnapshot();
+
+			/*
+			 * When the destination is DestTuplestore, the query snapshot is
+			 * required to read the result.  Since the snapshot is used as
+			 * holdSnapshot of the upper portal, its lifetime must be equal to
+			 * or longer than the life of the upper portal. If the queryDesc
+			 * comes from GetPortalByName, the lifetime of the snapshot should
+			 * be longer. On the other hand if the query comes from
+			 * GetCachedPlan, the snapshot would be freed before the upper
+			 * portal fetches the stored tuples. In that case, PortalRunUtility
+			 * should have set the active snapshot and everything goes fine
+			 * without setting the tuplestore's snapshot here.
+			 */
+			if (dest->mydest == DestTuplestore)
+			{
+				if (!ActiveSnapshotSet())
+					dest->snapshot = queryDesc->snapshot;
+				else
+					Assert(queryDesc->snapshot == GetActiveSnapshot());
+			}
 		}
 
 		if (!ScanDirectionIsNoMovement(direction))
@@ -965,6 +985,18 @@ PortalRunSelect(Portal portal,
 						portal->run_once);
 			nprocessed = queryDesc->estate->es_processed;
 			PopActiveSnapshot();
+
+			/*
+			 * Set query snapshot for the upper portal. See the comment above
+			 * for the details.
+			 */
+			if (dest->mydest == DestTuplestore)
+			{
+				if (!ActiveSnapshotSet())
+					dest->snapshot = queryDesc->snapshot;
+				else
+					Assert(queryDesc->snapshot == GetActiveSnapshot());
+			}
 		}
 
 		if (!ScanDirectionIsNoMovement(direction))
@@ -1038,6 +1070,16 @@ FillPortalStore(Portal portal, bool isTopLevel)
 			break;
 	}
 
+	/*
+	 * Set holdSnapshot if passed from the command. It should be set only if
+	 * holdSnapshot is not set yet.
+	 */
+	if (treceiver->snapshot)
+	{
+		Assert(portal->holdSnapshot == NULL);
+		portal->holdSnapshot = RegisterSnapshot(treceiver->snapshot);
+	}
+
 	/* Override portal completion data with actual command results */
 	if (qc.commandTag != CMDTAG_UNKNOWN)
 		CopyQueryCompletion(&portal->qc, &qc);
@@ -1063,9 +1105,22 @@ RunFromStore(Portal portal, ScanDirection direction, uint64 count,
 {
 	uint64		current_tuple_count = 0;
 	TupleTableSlot *slot;
+	bool		do_free_snapshot = false;
 
 	slot = MakeSingleTupleTableSlot(portal->tupDesc, &TTSOpsMinimalTuple);
 
+	/*
+	 * There's a case where the filler of the holdStore may set holdSnapshot
+	 * even if our caller didn't set one. Use that snapshot in that case.
+	 */
+	if (portal->portalSnapshot == NULL && portal->holdSnapshot)
+	{
+		PushActiveSnapshotWithLevel(portal->holdSnapshot,
+									portal->createLevel);
+		portal->portalSnapshot = GetActiveSnapshot();
+		do_free_snapshot = true;
+	}
+
 	dest->rStartup(dest, CMD_SELECT, portal->tupDesc);
 
 	if (ScanDirectionIsNoMovement(direction))
@@ -1114,6 +1169,14 @@ RunFromStore(Portal portal, ScanDirection direction, uint64 count,
 
 	dest->rShutdown(dest);
 
+	/* Delete the active snapshot if it is set in this function. */
+	if (do_free_snapshot)
+	{
+		Assert(portal->portalSnapshot == GetActiveSnapshot());
+		PopActiveSnapshot();
+		portal->portalSnapshot = NULL;
+	}
+
 	ExecDropSingleTupleTableSlot(slot);
 
 	return current_tuple_count;
diff --git a/src/include/tcop/dest.h b/src/include/tcop/dest.h
index 3c3eabae67..b90dbcc2f9 100644
--- a/src/include/tcop/dest.h
+++ b/src/include/tcop/dest.h
@@ -69,7 +69,7 @@
 
 #include "executor/tuptable.h"
 #include "tcop/cmdtag.h"
-
+#include "utils/snapshot.h"
 
 /* buffer size to use for command completion tags */
 #define COMPLETION_TAG_BUFSIZE	64
@@ -125,6 +125,8 @@ struct _DestReceiver
 	void		(*rShutdown) (DestReceiver *self);
 	/* Destroy the receiver object itself (if dynamically allocated) */
 	void		(*rDestroy) (DestReceiver *self);
+	/* tuplestore only: snapshot to use when reading the tuplestore */
+	Snapshot	snapshot;
 	/* CommandDest code for this receiver */
 	CommandDest mydest;
 	/* Private fields might appear beyond this point... */
diff --git a/src/test/regress/expected/portals.out b/src/test/regress/expected/portals.out
index 9da74876e1..98a7b04b4d 100644
--- a/src/test/regress/expected/portals.out
+++ b/src/test/regress/expected/portals.out
@@ -1536,3 +1536,12 @@ fetch backward all in c2;
 (3 rows)
 
 rollback;
+-- Check if snapshot is available for cursor while accessing toast value
+begin;
+create table t(a text);
+alter table t alter column a set storage external;
+insert into t values(repeat('x', 4096));
+declare c cursor for select a from t;
+-- the following shouldn't crash
+fetch all in c \gset result
+rollback;
diff --git a/src/test/regress/sql/portals.sql b/src/test/regress/sql/portals.sql
index eadf6ed942..545fad1e04 100644
--- a/src/test/regress/sql/portals.sql
+++ b/src/test/regress/sql/portals.sql
@@ -581,3 +581,13 @@ declare c2 scroll cursor for select generate_series(1,3) as g;
 fetch all in c2;
 fetch backward all in c2;
 rollback;
+
+-- Check if snapshot is available for cursor while accessing toast value
+begin;
+create table t(a text);
+alter table t alter column a set storage external;
+insert into t values(repeat('x', 4096));
+declare c cursor for select a from t;
+-- the following shouldn't crash
+fetch all in c \gset result
+rollback;
-- 
2.27.0

