From 35eddf8c663a41a6faa213edac6ce4086f18770d Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <pavel.stehule@gmail.com>
Date: Sun, 1 Jun 2025 21:20:16 +0200
Subject: [PATCH 10/15] svariableReceiver

allows to store result of the query to session variable

Check correct format of result - one column, one row.
---
 src/backend/executor/Makefile            |   1 +
 src/backend/executor/meson.build         |   1 +
 src/backend/executor/svariableReceiver.c | 172 +++++++++++++++++++++++
 src/backend/tcop/dest.c                  |   7 +
 src/include/executor/svariableReceiver.h |  22 +++
 src/include/tcop/dest.h                  |   1 +
 src/tools/pgindent/typedefs.list         |   1 +
 7 files changed, 205 insertions(+)
 create mode 100644 src/backend/executor/svariableReceiver.c
 create mode 100644 src/include/executor/svariableReceiver.h

diff --git a/src/backend/executor/Makefile b/src/backend/executor/Makefile
index 11118d0ce02..71248a34f26 100644
--- a/src/backend/executor/Makefile
+++ b/src/backend/executor/Makefile
@@ -76,6 +76,7 @@ OBJS = \
 	nodeWindowAgg.o \
 	nodeWorktablescan.o \
 	spi.o \
+	svariableReceiver.o \
 	tqueue.o \
 	tstoreReceiver.o
 
diff --git a/src/backend/executor/meson.build b/src/backend/executor/meson.build
index 2cea41f8771..491092fcc4c 100644
--- a/src/backend/executor/meson.build
+++ b/src/backend/executor/meson.build
@@ -64,6 +64,7 @@ backend_sources += files(
   'nodeWindowAgg.c',
   'nodeWorktablescan.c',
   'spi.c',
+  'svariableReceiver.c',
   'tqueue.c',
   'tstoreReceiver.c',
 )
diff --git a/src/backend/executor/svariableReceiver.c b/src/backend/executor/svariableReceiver.c
new file mode 100644
index 00000000000..41a967cb4b2
--- /dev/null
+++ b/src/backend/executor/svariableReceiver.c
@@ -0,0 +1,172 @@
+/*-------------------------------------------------------------------------
+ *
+ * svariableReceiver.c
+ *	  An implementation of DestReceiver that stores the result value in
+ *	  a session variable.
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/backend/executor/svariableReceiver.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+#include "miscadmin.h"
+
+#include "access/detoast.h"
+#include "catalog/pg_variable.h"
+#include "commands/session_variable.h"
+#include "executor/svariableReceiver.h"
+#include "storage/lock.h"
+#include "utils/builtins.h"
+#include "utils/lsyscache.h"
+#include "utils/syscache.h"
+
+/*
+ * This DestReceiver is used by the LET command for storing the result to a
+ * session variable.  The result has to have only one tuple with only one
+ * non-deleted attribute.  The row counter (field "rows") is incremented
+ * after receiving a row, and an error is raised when there are no rows or
+ * there are more than one received rows.  A received tuple cannot to have
+ * deleted attributes.  The value is detoasted before storing it in the
+ * session variable.
+ */
+typedef struct
+{
+	DestReceiver pub;
+	Oid			varid;
+	bool		need_detoast;	/* do we need to detoast the attribute? */
+	int			rows;			/* row counter */
+} SVariableState;
+
+/*
+ * Prepare to receive tuples from executor.
+ */
+static void
+svariableStartupReceiver(DestReceiver *self, int operation, TupleDesc typeinfo)
+{
+	SVariableState *myState = (SVariableState *) self;
+	LOCKTAG		locktag PG_USED_FOR_ASSERTS_ONLY;
+	Form_pg_attribute attr;
+	Oid			typid PG_USED_FOR_ASSERTS_ONLY;
+	Oid			collid PG_USED_FOR_ASSERTS_ONLY;
+	int32		typmod PG_USED_FOR_ASSERTS_ONLY;
+
+	Assert(myState->pub.mydest == DestVariable);
+	Assert(OidIsValid(myState->varid));
+	Assert(SearchSysCacheExists1(VARIABLEOID, myState->varid));
+	Assert(typeinfo->natts == 1);
+
+#ifdef USE_ASSERT_CHECKING
+
+	SET_LOCKTAG_OBJECT(locktag,
+					   MyDatabaseId,
+					   VariableRelationId,
+					   myState->varid,
+					   0);
+
+	Assert(LockHeldByMe(&locktag, AccessShareLock, false));
+
+#endif
+
+	attr = TupleDescAttr(typeinfo, 0);
+
+	Assert(!attr->attisdropped);
+
+#ifdef USE_ASSERT_CHECKING
+
+	get_session_variable_type_typmod_collid(myState->varid,
+											&typid,
+											&typmod,
+											&collid);
+
+	Assert(attr->atttypid == typid);
+	Assert(attr->atttypmod < 0 || attr->atttypmod == typmod);
+
+#endif
+
+	myState->need_detoast = attr->attlen == -1;
+	myState->rows = 0;
+}
+
+/*
+ * Receive a tuple from the executor and store it in the session variable.
+ */
+static bool
+svariableReceiveSlot(TupleTableSlot *slot, DestReceiver *self)
+{
+	SVariableState *myState = (SVariableState *) self;
+	Datum		value;
+	bool		isnull;
+	bool		freeval = false;
+
+	/* make sure the tuple is fully deconstructed */
+	slot_getallattrs(slot);
+
+	value = slot->tts_values[0];
+	isnull = slot->tts_isnull[0];
+
+	if (myState->need_detoast && !isnull && VARATT_IS_EXTERNAL(DatumGetPointer(value)))
+	{
+		value = PointerGetDatum(detoast_external_attr((struct varlena *)
+													  DatumGetPointer(value)));
+		freeval = true;
+	}
+
+	myState->rows += 1;
+
+	if (myState->rows > 1)
+		ereport(ERROR,
+				(errcode(ERRCODE_TOO_MANY_ROWS),
+				 errmsg("expression returned more than one row")));
+
+	SetSessionVariable(myState->varid, value, isnull);
+
+	if (freeval)
+		pfree(DatumGetPointer(value));
+
+	return true;
+}
+
+/*
+ * Clean up at end of the executor run
+ */
+static void
+svariableShutdownReceiver(DestReceiver *self)
+{
+	if (((SVariableState *) self)->rows == 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_NO_DATA_FOUND),
+				 errmsg("expression returned no rows")));
+}
+
+/*
+ * Destroy the receiver when we are done with it
+ */
+static void
+svariableDestroyReceiver(DestReceiver *self)
+{
+	pfree(self);
+}
+
+/*
+ * Initially create a DestReceiver object.
+ */
+DestReceiver *
+CreateVariableDestReceiver(Oid varid)
+{
+	SVariableState *self = (SVariableState *) palloc0(sizeof(SVariableState));
+
+	self->pub.receiveSlot = svariableReceiveSlot;
+	self->pub.rStartup = svariableStartupReceiver;
+	self->pub.rShutdown = svariableShutdownReceiver;
+	self->pub.rDestroy = svariableDestroyReceiver;
+	self->pub.mydest = DestVariable;
+
+	self->varid = varid;
+
+	return (DestReceiver *) self;
+}
diff --git a/src/backend/tcop/dest.c b/src/backend/tcop/dest.c
index b620766c938..b2f764b657f 100644
--- a/src/backend/tcop/dest.c
+++ b/src/backend/tcop/dest.c
@@ -38,6 +38,7 @@
 #include "executor/functions.h"
 #include "executor/tqueue.h"
 #include "executor/tstoreReceiver.h"
+#include "executor/svariableReceiver.h"
 #include "libpq/libpq.h"
 #include "libpq/pqformat.h"
 
@@ -155,6 +156,9 @@ CreateDestReceiver(CommandDest dest)
 
 		case DestExplainSerialize:
 			return CreateExplainSerializeDestReceiver(NULL);
+
+		case DestVariable:
+			return CreateVariableDestReceiver(InvalidOid);
 	}
 
 	/* should never get here */
@@ -191,6 +195,7 @@ EndCommand(const QueryCompletion *qc, CommandDest dest, bool force_undecorated_o
 		case DestTransientRel:
 		case DestTupleQueue:
 		case DestExplainSerialize:
+		case DestVariable:
 			break;
 	}
 }
@@ -237,6 +242,7 @@ NullCommand(CommandDest dest)
 		case DestTransientRel:
 		case DestTupleQueue:
 		case DestExplainSerialize:
+		case DestVariable:
 			break;
 	}
 }
@@ -281,6 +287,7 @@ ReadyForQuery(CommandDest dest)
 		case DestTransientRel:
 		case DestTupleQueue:
 		case DestExplainSerialize:
+		case DestVariable:
 			break;
 	}
 }
diff --git a/src/include/executor/svariableReceiver.h b/src/include/executor/svariableReceiver.h
new file mode 100644
index 00000000000..db44d8b94c6
--- /dev/null
+++ b/src/include/executor/svariableReceiver.h
@@ -0,0 +1,22 @@
+/*-------------------------------------------------------------------------
+ *
+ * svariableReceiver.h
+ *	  prototypes for svariableReceiver.c
+ *
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/executor/svariableReceiver.h
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef SVARIABLE_RECEIVER_H
+#define SVARIABLE_RECEIVER_H
+
+#include "tcop/dest.h"
+
+extern DestReceiver *CreateVariableDestReceiver(Oid varid);
+
+#endif							/* SVARIABLE_RECEIVER_H */
diff --git a/src/include/tcop/dest.h b/src/include/tcop/dest.h
index 00c092e3d7c..6ce3ea0e617 100644
--- a/src/include/tcop/dest.h
+++ b/src/include/tcop/dest.h
@@ -97,6 +97,7 @@ typedef enum
 	DestTransientRel,			/* results sent to transient relation */
 	DestTupleQueue,				/* results sent to tuple queue */
 	DestExplainSerialize,		/* results are serialized and discarded */
+	DestVariable,				/* results sent to session variable */
 } CommandDest;
 
 /* ----------------
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 4394dc3f16b..77f864804ad 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2651,6 +2651,7 @@ STRLEN
 SV
 SVariableData
 SVariable
+SVariableState
 SYNCHRONIZATION_BARRIER
 SYSTEM_INFO
 SampleScan
-- 
2.50.0

