v20260403-2-0005-svariableReceiver.patch

From 654e7eb78cb07434aebfb48096a66b22e3c46094 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 | 139 +++++++++++++++++++++++
 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, 225 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 19c153dae3f..cd171a68cbf 100644
--- a/src/backend/commands/session_variable.c
+++ b/src/backend/commands/session_variable.c
@@ -167,6 +167,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 dc45be0b2ce..a572b6dab7c 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..f8fbb7a8e71
--- /dev/null
+++ b/src/backend/executor/svariableReceiver.c
@@ -0,0 +1,139 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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 "access/htup_details.h"
+#include "commands/session_variable.h"
+#include "executor/svariableReceiver.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 assignment to session variable have to be postponed until we are
+ * sure so only one row was received.
+ */
+typedef struct
+{
+	DestReceiver pub;
+	char	   *varname;
+	int			rows;			/* row counter */
+	TupleDesc	tupdesc;
+	HeapTuple	tuple;
+	MemoryContext tuple_cxt;	/* holds a value before storing to variable */
+} SVariableState;
+
+/*
+ * Prepare to receive tuples from executor.
+ */
+static void
+svariableStartupReceiver(DestReceiver *self, int operation, TupleDesc typeinfo)
+{
+	SVariableState *myState = (SVariableState *) self;
+
+	Assert(myState->pub.mydest == DestVariable);
+	Assert(typeinfo->natts == 1);
+
+	myState->rows = 0;
+	myState->tupdesc = typeinfo;
+	myState->tuple = NULL;
+	myState->tuple_cxt = CurrentMemoryContext;
+}
+
+/*
+ * Receive a tuple from the executor and store it in the buffer
+ */
+static bool
+svariableReceiveSlot(TupleTableSlot *slot, DestReceiver *self)
+{
+	SVariableState *myState = (SVariableState *) self;
+	MemoryContext oldcxt;
+
+	if (++myState->rows > 1)
+		ereport(ERROR,
+				(errcode(ERRCODE_TOO_MANY_ROWS),
+				 errmsg("expression returned more than one row")));
+
+	/*
+	 * We cannot to assign received value directly, so we should to
+	 * save received value in the buffer.
+	 */
+	oldcxt = MemoryContextSwitchTo(myState->tuple_cxt);
+	myState->tuple = ExecCopySlotHeapTuple(slot);
+	MemoryContextSwitchTo(oldcxt);
+
+	return true;
+}
+
+/*
+ * Clean up at end of the executor run
+ */
+static void
+svariableShutdownReceiver(DestReceiver *self)
+{
+	SVariableState *myState = (SVariableState *) self;
+	Form_pg_attribute attr;
+	Datum		value;
+	bool		isnull;
+
+	if (myState->rows == 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_NO_DATA_FOUND),
+				 errmsg("expression returned no rows")));
+
+	attr = TupleDescAttr(myState->tupdesc, 0);
+	Assert(!attr->attisdropped);
+
+	value = heap_getattr(myState->tuple, 1, myState->tupdesc, &isnull);
+
+	SetSessionVariableWithTypecheck(myState->varname,
+									attr->atttypid, attr->atttypmod,
+									value, isnull);
+
+	heap_freetuple(myState->tuple);
+}
+
+/*
+ * Destroy the receiver when we are done with it
+ */
+static void
+svariableDestroyReceiver(DestReceiver *self)
+{
+	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 = varname;
+
+	return (DestReceiver *) self;
+}
diff --git a/src/backend/tcop/dest.c b/src/backend/tcop/dest.c
index fb163930c89..13fe536b432 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 */
@@ -192,6 +196,7 @@ EndCommand(const QueryCompletion *qc, CommandDest dest, bool force_undecorated_o
 		case DestTransientRel:
 		case DestTupleQueue:
 		case DestExplainSerialize:
+		case DestVariable:
 			break;
 	}
 }
@@ -238,6 +243,7 @@ NullCommand(CommandDest dest)
 		case DestTransientRel:
 		case DestTupleQueue:
 		case DestExplainSerialize:
+		case DestVariable:
 			break;
 	}
 }
@@ -282,6 +288,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 4e4f532d8cc..1e7043dc7f7 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 04f1bd17dcc..22198dfafef 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2747,6 +2747,7 @@ STRLEN
 SV
 SVariableData
 SVariable
+SVariableState
 SYNCHRONIZATION_BARRIER
 SYSTEM_INFO
 SampleScan
-- 
2.53.0