From f8b0cfac9c4a948b5106a967777478affa2c5328 Mon Sep 17 00:00:00 2001
From: Laurenz Albe <laurenz.albe@cybertec.at>
Date: Wed, 13 Nov 2024 15:02:13 +0100
Subject: [PATCH 13/22] Implementation ON TRANSACTION END RESET clause

This is simple patch - just add special flag to session variable memory entry.
The entries with active this flag are removed from sessionvars hash table
at transaction end. The "TRANSACTION END" is synonyms for "COMMIT ROLLBACK"
but the "TRANSACTION END" is more illustrative and less confusing then "COMMIT ROLLBACK".
---
 doc/src/sgml/catalogs.sgml                    |  3 +-
 doc/src/sgml/ref/create_variable.sgml         | 13 ++++-
 src/backend/commands/session_variable.c       | 51 +++++++++++++++++++
 src/backend/parser/gram.y                     |  1 +
 src/bin/pg_dump/pg_dump.c                     | 11 ++++
 src/bin/pg_dump/pg_dump.h                     |  1 +
 src/bin/pg_dump/t/002_pg_dump.pl              | 18 +++++++
 src/bin/psql/describe.c                       |  1 +
 src/include/catalog/pg_variable.h             |  1 +
 .../isolation/expected/session-variable.out   |  4 +-
 .../isolation/specs/session-variable.spec     |  4 +-
 .../regress/expected/session_variables.out    | 34 +++++++++++++
 src/test/regress/sql/session_variables.sql    | 20 ++++++++
 13 files changed, 158 insertions(+), 4 deletions(-)

diff --git a/doc/src/sgml/catalogs.sgml b/doc/src/sgml/catalogs.sgml
index f651ad90b10..20fc5d60cf5 100644
--- a/doc/src/sgml/catalogs.sgml
+++ b/doc/src/sgml/catalogs.sgml
@@ -9845,7 +9845,8 @@ SCRAM-SHA-256$<replaceable>&lt;iteration count&gt;</replaceable>:<replaceable>&l
       </para>
       <para>
        Action performed at end of transaction:
-       <literal>n</literal> = no action, <literal>d</literal> = drop the variable.
+       <literal>n</literal> = no action, <literal>d</literal> = drop the variable,
+       <literal>r</literal> = reset the variable to its default value.
       </para></entry>
      </row>
 
diff --git a/doc/src/sgml/ref/create_variable.sgml b/doc/src/sgml/ref/create_variable.sgml
index 81ffbcc3c62..d813668a776 100644
--- a/doc/src/sgml/ref/create_variable.sgml
+++ b/doc/src/sgml/ref/create_variable.sgml
@@ -27,7 +27,7 @@ PostgreSQL documentation
  <refsynopsisdiv>
 <synopsis>
 CREATE [ { TEMPORARY | TEMP } ] VARIABLE [ IF NOT EXISTS ] <replaceable class="parameter">name</replaceable> [ AS ] <replaceable class="parameter">data_type</replaceable> [ COLLATE <replaceable class="parameter">collation</replaceable> ]
-    [ ON COMMIT DROP ]
+    [ { ON COMMIT DROP | ON TRANSACTION END RESET } ]
 </synopsis>
  </refsynopsisdiv>
  <refsect1>
@@ -117,6 +117,17 @@ CREATE [ { TEMPORARY | TEMP } ] VARIABLE [ IF NOT EXISTS ] <replaceable class="p
     </listitem>
    </varlistentry>
 
+   <varlistentry id="sql-createvariable-on-transaction-end-reset">
+    <term><literal>ON TRANSACTION END RESET</literal></term>
+    <listitem>
+     <para>
+      The <literal>ON TRANSACTION END RESET</literal> clause causes the session
+      variable to be reset to its default value when the transaction is committed
+      or rolled back.
+     </para>
+    </listitem>
+   </varlistentry>
+
   </variablelist>
  </refsect1>
 
diff --git a/src/backend/commands/session_variable.c b/src/backend/commands/session_variable.c
index 5ea20fb0cf8..f16c31bf3a9 100644
--- a/src/backend/commands/session_variable.c
+++ b/src/backend/commands/session_variable.c
@@ -91,6 +91,8 @@ typedef struct SVariableData
 	void	   *domain_check_extra;
 	LocalTransactionId domain_check_extra_lxid;
 
+	bool		reset_at_eox;
+
 	/*
 	 * Top level local transaction id of the last transaction that dropped the
 	 * variable, if any.  We need this information to avoid freeing memory for
@@ -118,6 +120,12 @@ static MemoryContext SVariableMemoryContext = NULL;
 /* becomes true when we receive a sinval message */
 static bool needs_validation = false;
 
+/*
+ * true, when some used session variable has ON COMMIT DROP
+ * or ON TRANSACTION END RESET clauses
+ */
+static bool has_session_variables_with_reset_at_eox = false;
+
 /*
  * The content of dropped session variables is not removed immediately.  If
  * possible, we do that at the end of the transaction.  But we cannot do that
@@ -393,6 +401,32 @@ remove_invalid_session_variables(bool atEOX)
 	}
 }
 
+/*
+ * remove entries marked as "reset_at_eox"
+ */
+static void
+remove_session_variables_with_reset_at_eox(void)
+{
+	HASH_SEQ_STATUS status;
+	SVariable	svar;
+
+	if (!sessionvars)
+		return;
+
+	/* leave quckly, when there are not that variables */
+	if (!has_session_variables_with_reset_at_eox)
+		return;
+
+	hash_seq_init(&status, sessionvars);
+	while ((svar = (SVariable) hash_seq_search(&status)) != NULL)
+	{
+		if (svar->reset_at_eox)
+			hash_search(sessionvars, &svar->varid, HASH_REMOVE, NULL);
+	}
+
+	has_session_variables_with_reset_at_eox = false;
+}
+
 /*
   * Perform ON COMMIT DROP for temporary session variables,
   * and remove all dropped variables from memory.
@@ -400,6 +434,8 @@ remove_invalid_session_variables(bool atEOX)
 void
 AtPreEOXact_SessionVariables(bool isCommit)
 {
+	remove_session_variables_with_reset_at_eox();
+
 	if (isCommit)
 	{
 		if (xact_drop_items)
@@ -511,6 +547,21 @@ setup_session_variable(SVariable svar, Oid varid)
 	svar->domain_check_extra = NULL;
 	svar->domain_check_extra_lxid = InvalidLocalTransactionId;
 
+	/*
+	 * We don't need to explicitly reset variables marked ON COMMIT DROP. It
+	 * can be done by sinval message processing. But this processing can be
+	 * postponed due aborted transaction. On second hand there is not a
+	 * reason, why don't do it at transaction end immediately.
+	 */
+	if (varform->varxactendaction == VARIABLE_XACTEND_RESET ||
+		varform->varxactendaction == VARIABLE_XACTEND_DROP)
+	{
+		svar->reset_at_eox = true;
+		has_session_variables_with_reset_at_eox = true;
+	}
+	else
+		svar->reset_at_eox = false;
+
 	svar->drop_lxid = InvalidTransactionId;
 
 	svar->isnull = true;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 7c089549b4b..83cf91d3efb 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -5258,6 +5258,7 @@ CreateSessionVarStmt:
  * transaction end like tables.
  */
 XactEndActionOption:  ON COMMIT DROP				{ $$ = VARIABLE_XACTEND_DROP; }
+			| ON TRANSACTION END_P RESET			{ $$ = VARIABLE_XACTEND_RESET; }
 			| /*EMPTY*/								{ $$ = VARIABLE_XACTEND_NOOP; }
 		;
 
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 9e32c478412..f2f3779c672 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -5436,6 +5436,7 @@ getVariables(Archive *fout)
 	int			i_varnamespace;
 	int			i_vartype;
 	int			i_vartypname;
+	int			i_varxactendaction;
 	int			i_varowner;
 	int			i_varcollation;
 	int			i_varacl;
@@ -5453,6 +5454,7 @@ getVariables(Archive *fout)
 	/* get the variables in current database */
 	appendPQExpBuffer(query,
 					  "SELECT v.tableoid, v.oid, v.varname,\n"
+					  "       v.varxactendaction,\n"
 					  "       v.varnamespace, v.vartype,\n"
 					  "       pg_catalog.format_type(v.vartype, v.vartypmod) as vartypname,\n"
 					  "       CASE WHEN v.varcollation <> t.typcollation "
@@ -5475,6 +5477,7 @@ getVariables(Archive *fout)
 	i_varnamespace = PQfnumber(res, "varnamespace");
 	i_vartype = PQfnumber(res, "vartype");
 	i_vartypname = PQfnumber(res, "vartypname");
+	i_varxactendaction = PQfnumber(res, "varxactendaction");
 	i_varcollation = PQfnumber(res, "varcollation");
 
 	i_varowner = PQfnumber(res, "varowner");
@@ -5498,6 +5501,9 @@ getVariables(Archive *fout)
 
 		varinfo[i].vartype = atooid(PQgetvalue(res, i, i_vartype));
 		varinfo[i].vartypname = pg_strdup(PQgetvalue(res, i, i_vartypname));
+		varinfo[i].varxactendaction =
+			pg_strdup(PQgetvalue(res, i, i_varxactendaction));
+
 		varinfo[i].varcollation = atooid(PQgetvalue(res, i, i_varcollation));
 
 		varinfo[i].dacl.acl = pg_strdup(PQgetvalue(res, i, i_varacl));
@@ -5538,6 +5544,7 @@ dumpVariable(Archive *fout, const VariableInfo *varinfo)
 	PQExpBuffer query;
 	char	   *qualvarname;
 	const char *vartypname;
+	const char *varxactendaction;
 	Oid			varcollation;
 
 	/* skip if not to be dumped */
@@ -5549,6 +5556,7 @@ dumpVariable(Archive *fout, const VariableInfo *varinfo)
 
 	qualvarname = pg_strdup(fmtQualifiedDumpable(varinfo));
 	vartypname = varinfo->vartypname;
+	varxactendaction = varinfo->varxactendaction;
 	varcollation = varinfo->varcollation;
 
 	appendPQExpBuffer(delq, "DROP VARIABLE %s;\n",
@@ -5567,6 +5575,9 @@ dumpVariable(Archive *fout, const VariableInfo *varinfo)
 							  fmtQualifiedDumpable(coll));
 	}
 
+	if (strcmp(varxactendaction, "r") == 0)
+		appendPQExpBuffer(query, " ON TRANSACTION END RESET");
+
 	appendPQExpBuffer(query, ";\n");
 
 	if (varinfo->dobj.dump & DUMP_COMPONENT_DEFINITION)
diff --git a/src/bin/pg_dump/pg_dump.h b/src/bin/pg_dump/pg_dump.h
index 8a62b5b4128..a9ac70d7616 100644
--- a/src/bin/pg_dump/pg_dump.h
+++ b/src/bin/pg_dump/pg_dump.h
@@ -717,6 +717,7 @@ typedef struct _VariableInfo
 	DumpableAcl dacl;
 	Oid			vartype;
 	char	   *vartypname;
+	char	   *varxactendaction;
 	char	   *varacl;
 	char	   *rvaracl;
 	char	   *initvaracl;
diff --git a/src/bin/pg_dump/t/002_pg_dump.pl b/src/bin/pg_dump/t/002_pg_dump.pl
index 6423e7b9d27..7b0cc2c5d94 100644
--- a/src/bin/pg_dump/t/002_pg_dump.pl
+++ b/src/bin/pg_dump/t/002_pg_dump.pl
@@ -3987,6 +3987,24 @@ my %tests = (
 		},
 	},
 
+	'CREATE VARIABLE test_variable ON TRANSACTION END RESET' => {
+		all_runs     => 1,
+		catch_all    => 'CREATE ... commands',
+		create_order => 61,
+		create_sql   => 'CREATE VARIABLE dump_test.variable2 AS integer ON TRANSACTION END RESET;',
+		regexp => qr/^
+			\QCREATE VARIABLE dump_test.variable2 AS integer ON TRANSACTION END RESET;\E/xm,
+		like => {
+			%full_runs,
+			%dump_test_schema_runs,
+			section_pre_data => 1,
+		},
+		unlike => {
+			exclude_dump_test_schema => 1,
+			only_dump_measurement    => 1,
+		},
+	},
+
 	'CREATE VIEW test_view' => {
 		create_order => 61,
 		create_sql => 'CREATE VIEW dump_test.test_view
diff --git a/src/bin/psql/describe.c b/src/bin/psql/describe.c
index d6d30e73ee8..3f743676ca7 100644
--- a/src/bin/psql/describe.c
+++ b/src/bin/psql/describe.c
@@ -5273,6 +5273,7 @@ listVariables(const char *pattern, bool verbose)
 					  "  pg_catalog.pg_get_userbyid(v.varowner) as \"%s\",\n"
 					  "  CASE v.varxactendaction\n"
 					  "    WHEN 'd' THEN 'ON COMMIT DROP'\n"
+					  "    WHEN 'r' THEN 'ON TRANSACTION END RESET'\n"
 					  "  END as \"%s\"\n",
 					  gettext_noop("Schema"),
 					  gettext_noop("Name"),
diff --git a/src/include/catalog/pg_variable.h b/src/include/catalog/pg_variable.h
index 0fc672bc2ba..008f2ee7c17 100644
--- a/src/include/catalog/pg_variable.h
+++ b/src/include/catalog/pg_variable.h
@@ -75,6 +75,7 @@ typedef enum VariableXactEndAction
 {
 	VARIABLE_XACTEND_NOOP = 'n',	/* NOOP */
 	VARIABLE_XACTEND_DROP = 'd',	/* ON COMMIT DROP */
+	VARIABLE_XACTEND_RESET = 'r',	/* ON TRANSACTION END RESET */
 }			VariableXactEndAction;
 
 /* ----------------
diff --git a/src/test/isolation/expected/session-variable.out b/src/test/isolation/expected/session-variable.out
index a609797dc5d..e9a254bcd12 100644
--- a/src/test/isolation/expected/session-variable.out
+++ b/src/test/isolation/expected/session-variable.out
@@ -86,11 +86,12 @@ myvar
 
 step sr1: ROLLBACK;
 
-starting permutation: create3 let3 s3 o_c_d create4 let4 drop4 drop3 inval3 discard sc3 state
+starting permutation: create3 let3 s3 o_c_d o_eox_r create4 let4 drop4 drop3 inval3 discard sc3 clean state
 step create3: CREATE VARIABLE myvar3 AS text;
 step let3: LET myvar3 = 'test';
 step s3: BEGIN;
 step o_c_d: CREATE TEMP VARIABLE myvar_o_c_d AS text ON COMMIT DROP;
+step o_eox_r: CREATE VARIABLE myvar_o_eox_r AS text ON TRANSACTION END RESET; LET myvar_o_eox_r = 'test';
 step create4: CREATE VARIABLE myvar4 AS text;
 step let4: LET myvar4 = 'test';
 step drop4: DROP VARIABLE myvar4;
@@ -103,6 +104,7 @@ t
 
 step discard: DISCARD VARIABLES;
 step sc3: COMMIT;
+step clean: DROP VARIABLE myvar_o_eox_r;
 step state: SELECT varname FROM pg_variable;
 varname
 -------
diff --git a/src/test/isolation/specs/session-variable.spec b/src/test/isolation/specs/session-variable.spec
index 45e65d4085d..5d089c8a4ed 100644
--- a/src/test/isolation/specs/session-variable.spec
+++ b/src/test/isolation/specs/session-variable.spec
@@ -25,12 +25,14 @@ session s3
 step s3			{ BEGIN; }
 step let3		{ LET myvar3 = 'test'; }
 step o_c_d		{ CREATE TEMP VARIABLE myvar_o_c_d AS text ON COMMIT DROP; }
+step o_eox_r	{ CREATE VARIABLE myvar_o_eox_r AS text ON TRANSACTION END RESET; LET myvar_o_eox_r = 'test'; }
 step create4	{ CREATE VARIABLE myvar4 AS text; }
 step let4		{ LET myvar4 = 'test'; }
 step drop4		{ DROP VARIABLE myvar4; }
 step inval3		{ SELECT COUNT(*) >= 0 FROM pg_foreign_table; }
 step discard	{ DISCARD VARIABLES; }
 step sc3		{ COMMIT; }
+step clean		{ DROP VARIABLE myvar_o_eox_r; }
 step state		{ SELECT varname FROM pg_variable; }
 
 session s4
@@ -48,4 +50,4 @@ permutation let val dbg drop create dbg val
 # calling the dbg step after the concurrent drop
 permutation let val s1 dbg drop create dbg val sr1
 # test for DISCARD ALL when all internal queues have actions registered
-permutation create3 let3 s3 o_c_d create4 let4 drop4 drop3 inval3 discard sc3 state
+permutation create3 let3 s3 o_c_d o_eox_r create4 let4 drop4 drop3 inval3 discard sc3 clean state
diff --git a/src/test/regress/expected/session_variables.out b/src/test/regress/expected/session_variables.out
index 0d175cb756b..37669e6ba1f 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -2323,3 +2323,37 @@ SELECT count(*) FROM pg_session_variables();
      0
 (1 row)
 
+CREATE VARIABLE var1 AS int ON TRANSACTION END RESET;
+BEGIN;
+  LET var1 = 100;
+  SELECT var1;
+ var1 
+------
+  100
+(1 row)
+
+COMMIT;
+-- should be NULL;
+SELECT var1 IS NULL;
+ ?column? 
+----------
+ t
+(1 row)
+
+BEGIN;
+  LET var1 = 100;
+  SELECT var1;
+ var1 
+------
+  100
+(1 row)
+
+ROLLBACK;
+-- should be NULL
+SELECT var1 IS NULL;
+ ?column? 
+----------
+ t
+(1 row)
+
+DROP VARIABLE var1;
diff --git a/src/test/regress/sql/session_variables.sql b/src/test/regress/sql/session_variables.sql
index 2c62ee1f7cd..16274b58a0d 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -1593,3 +1593,23 @@ COMMIT;
 SELECT count(*) FROM pg_variable WHERE varname = 'var1';
 -- should be zero
 SELECT count(*) FROM pg_session_variables();
+
+CREATE VARIABLE var1 AS int ON TRANSACTION END RESET;
+
+BEGIN;
+  LET var1 = 100;
+  SELECT var1;
+COMMIT;
+
+-- should be NULL;
+SELECT var1 IS NULL;
+
+BEGIN;
+  LET var1 = 100;
+  SELECT var1;
+ROLLBACK;
+
+-- should be NULL
+SELECT var1 IS NULL;
+
+DROP VARIABLE var1;
-- 
2.47.1

