From 27ee61a899fe15ccfa0629ded0ddc8001781bcd0 Mon Sep 17 00:00:00 2001 From: "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