From 4e486f22b90cdad8c988b0b8204aff5976d5d937 Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <okbob@github.com>
Date: Sun, 1 Jun 2025 21:20:16 +0200
Subject: [PATCH 05/11] svariableReceiver

allows to store result of the query to session variable

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

diff --git a/src/backend/commands/session_variable.c b/src/backend/commands/session_variable.c
index 201bc048d8b..66ea1ec2eba 100644
--- a/src/backend/commands/session_variable.c
+++ b/src/backend/commands/session_variable.c
@@ -171,6 +171,56 @@ GetSessionVariableWithTypecheck(char *varname,
 	return result;
 }
 
+/*
+ * Store the given value in a session variable in the cache.
+ */
+void
+SetSessionVariableWithTypecheck(char *varname,
+								Oid typid, int32 typmod,
+								Datum value, bool isnull)
+{
+	SVariable	svar;
+
+	svar = search_variable(varname);
+
+	if (svar->vartype != typid || svar->vartypmod != typmod)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+				 errmsg("session variable %s is not of a type %s but type %s",
+						varname,
+						format_type_with_typemod(typid, typmod),
+						format_type_with_typemod(svar->vartype, svar->vartypmod))));
+
+	/* only owner can set content of variable */
+	if (svar->varowner != GetUserId() && !superuser())
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("permission denied for session variable %s",
+						varname)));
+
+	if (!svar->typbyval)
+	{
+		if (!isnull)
+		{
+			MemoryContext oldcxt;
+
+			/*
+			 * Do copy of value in session variables context. This operation
+			 * can fail, so do it before releasing the old content.
+			 */
+			oldcxt = MemoryContextSwitchTo(SVariableMemoryContext);
+			value = datumCopy(value, svar->typbyval, svar->typlen);
+			MemoryContextSwitchTo(oldcxt);
+		}
+
+		if (!svar->isnull)
+			pfree(DatumGetPointer(svar->value));
+	}
+
+	svar->value = value;
+	svar->isnull = isnull;
+}
+
 /*
  * Creates a new variable - does new entry in sessionvars
  *
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..b2709e9211b
--- /dev/null
+++ b/src/backend/executor/svariableReceiver.c
@@ -0,0 +1,149 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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 "access/htup_details.h"
+#include "commands/session_variable.h"
+#include "executor/svariableReceiver.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;
+	char	   *varname;
+	Oid			typid;
+	int32		typmod;
+	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;
+	Form_pg_attribute attr;
+
+	Assert(myState->pub.mydest == DestVariable);
+	Assert(typeinfo->natts == 1);
+
+	attr = TupleDescAttr(typeinfo, 0);
+
+	Assert(!attr->attisdropped);
+
+	myState->typid = attr->atttypid;
+	myState->typmod = attr->atttypmod;
+
+	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")));
+
+	SetSessionVariableWithTypecheck(myState->varname,
+									myState->typid, myState->typmod,
+									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(((SVariableState *) self)->varname);
+	pfree(self);
+}
+
+/*
+ * Initially create a DestReceiver object.
+ */
+DestReceiver *
+CreateVariableDestReceiver(char *varname)
+{
+	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->varname = pstrdup(varname);
+
+	return (DestReceiver *) self;
+}
diff --git a/src/backend/tcop/dest.c b/src/backend/tcop/dest.c
index b620766c938..0c1eeeb22a6 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(NULL);
 	}
 
 	/* 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/commands/session_variable.h b/src/include/commands/session_variable.h
index 3687490bcb1..610b757899e 100644
--- a/src/include/commands/session_variable.h
+++ b/src/include/commands/session_variable.h
@@ -23,6 +23,9 @@ extern void CreateVariable(ParseState *pstate, CreateSessionVarStmt *stmt);
 extern void DropVariableByName(char *varname);
 
 extern Datum GetSessionVariableWithTypecheck(char *varname, Oid typid, int32 typmod, bool *isnull);
+extern void SetSessionVariableWithTypecheck(char *varname,
+											Oid typid, int32 typmod,
+											Datum value, bool isnull);
 
 extern void get_session_variable_type_typmod_collid(char *varname,
 													Oid *typid,
diff --git a/src/include/executor/svariableReceiver.h b/src/include/executor/svariableReceiver.h
new file mode 100644
index 00000000000..dd01c93c9e8
--- /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(char *varname);
+
+#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 40b920d91e3..23ff5dc2281 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2671,6 +2671,7 @@ STRLEN
 SV
 SVariableData
 SVariable
+SVariableState
 SYNCHRONIZATION_BARRIER
 SYSTEM_INFO
 SampleScan
-- 
2.52.0

