v9-0008-psql-Add-parameterset-meta-command.patch

From 851c450daa20662bbdbd9d4b3fe4bb9e6584f23a Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Mon, 15 Jan 2024 13:01:21 +0100
Subject: [PATCH v9 08/13] psql: Add \parameterset meta-command

To be able to easily test the new ParameterSet protocol message this
introduces a new meta-command to psql: \parameterset

An additional usecase for this meta-command is using it to configure a
connection to connection poolers from the client side.
---
 doc/src/sgml/ref/psql-ref.sgml     | 13 +++++++++
 src/bin/psql/command.c             | 45 ++++++++++++++++++++++++++++++
 src/bin/psql/common.c              | 10 +++++++
 src/bin/psql/settings.h            |  2 ++
 src/test/regress/expected/psql.out | 26 +++++++++++++++++
 src/test/regress/sql/psql.sql      | 14 ++++++++++
 6 files changed, 110 insertions(+)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index cc7d797159b..4954399874c 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -2779,6 +2779,19 @@ lo_import 152801
         </listitem>
       </varlistentry>
 
+      <varlistentry id="app-psql-meta-command-parameterset">
+        <term><literal>\parameterset <replaceable class="parameter">name</replaceable> <replaceable class="parameter">value</replaceable></literal></term>
+        <listitem>
+        <para>
+         This command can be used to change a server parameter using the
+         ParameterSet protocol message. This has the effect as using an
+         equivalent <literal>SET</literal> command. The main difference is that
+         Some connection poolers can only intercept parameter changes that are
+         changed using the ParameterSet message.
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry id="app-psql-meta-command-password">
         <term><literal>\password [ <replaceable class="parameter">username</replaceable> ]</literal></term>
         <listitem>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 5c906e48068..f4ba7f8d5d0 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -116,6 +116,7 @@ static backslashResult exec_command_lo(PsqlScanState scan_state, bool active_bra
 static backslashResult exec_command_out(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_print(PsqlScanState scan_state, bool active_branch,
 										  PQExpBuffer query_buf, PQExpBuffer previous_buf);
+static backslashResult exec_command_parameterset(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_password(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_prompt(PsqlScanState scan_state, bool active_branch,
 										   const char *cmd);
@@ -378,6 +379,8 @@ exec_command(const char *cmd,
 	else if (strcmp(cmd, "p") == 0 || strcmp(cmd, "print") == 0)
 		status = exec_command_print(scan_state, active_branch,
 									query_buf, previous_buf);
+	else if (strcmp(cmd, "parameterset") == 0)
+		status = exec_command_parameterset(scan_state, active_branch);
 	else if (strcmp(cmd, "password") == 0)
 		status = exec_command_password(scan_state, active_branch);
 	else if (strcmp(cmd, "prompt") == 0)
@@ -2095,6 +2098,48 @@ exec_command_print(PsqlScanState scan_state, bool active_branch,
 	return PSQL_CMD_SKIP_LINE;
 }
 
+/*
+ * \parameterset -- set query parameters
+ */
+static backslashResult
+exec_command_parameterset(PsqlScanState scan_state, bool active_branch)
+{
+	if (!active_branch)
+	{
+		ignore_slash_options(scan_state);
+		return PSQL_CMD_SKIP_LINE;
+	}
+
+	for (int i = 0; i < lengthof(pset.parameterset_args); i++)
+		pset.parameterset_args[i] = psql_scan_slash_option(scan_state,
+														   OT_NORMAL,
+														   NULL,
+														   true);
+
+	if (pset.parameterset_args[1] == NULL)
+	{
+		pg_log_error("\\parameterset needs two arguments");
+		goto error;
+	}
+
+	if (strncmp("_pq_.", pset.parameterset_args[0], 5) == 0)
+	{
+		pg_log_error("\\parameterset cannot be used to change protocol extensions parameters");
+		goto error;
+	}
+
+	pset.parameterset_flag = true;
+
+	return PSQL_CMD_SEND;
+
+error:
+	free(pset.parameterset_args[0]);
+	free(pset.parameterset_args[1]);
+	pset.parameterset_args[0] = NULL;
+	pset.parameterset_args[1] = NULL;
+	return PSQL_CMD_ERROR;
+}
+
 /*
  * \password -- set user password
  */
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 76e01b02a39..a2d223c7fd6 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1272,6 +1272,14 @@ sendquery_cleanup:
 		pset.ctv_args[i] = NULL;
 	}
 
+	/* reset \parameterset trigger */
+	pset.parameterset_flag = false;
+	for (i = 0; i < lengthof(pset.parameterset_args); i++)
+	{
+		pg_free(pset.parameterset_args[i]);
+		pset.parameterset_args[i] = NULL;
+	}
+
 	return OK;
 }
 
@@ -1435,6 +1443,8 @@ ExecQueryAndProcessResults(const char *query,
 
 	if (pset.bind_flag)
 		success = PQsendQueryParams(pset.db, query, pset.bind_nparams, NULL, (const char *const *) pset.bind_params, NULL, NULL, 0);
+	else if (pset.parameterset_flag)
+		success = PQsendParameterSet(pset.db, pset.parameterset_args[0], pset.parameterset_args[1]);
 	else
 		success = PQsendQuery(pset.db, query);
 
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 505f99d8e47..c93f8792443 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -105,6 +105,8 @@ typedef struct _psqlSettings
 
 	bool		notty;			/* stdin or stdout is not a tty (as determined
 								 * on startup) */
+	bool		parameterset_flag;	/* one-shot request to crosstab result */
+	char	   *parameterset_args[2];	/* \parameterset arguments */
 	enum trivalue getPassword;	/* prompt the user for a username and password */
 	FILE	   *cur_cmd_source; /* describe the status of the current main
 								 * loop */
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index ad027725624..46421325a58 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -129,6 +129,31 @@ ERROR:  cannot insert multiple commands into a prepared statement
 -- bind error
 SELECT $1, $2 \bind 'foo' \g
 ERROR:  bind message supplies 1 parameters, but prepared statement "" requires 2
+-- \parameterset (protocol version 3.1)
+\parameterset work_mem 1234
+SHOW work_mem;
+ work_mem 
+----------
+ 1234kB
+(1 row)
+
+\parameterset work_mem 'abc'
+ERROR:  invalid value for parameter "work_mem": "abc"
+RESET work_mem;
+\parameterset search_path '"some quoted string",and_one_without,""'
+SHOW search_path;
+               search_path               
+-----------------------------------------
+ "some quoted string",and_one_without,""
+(1 row)
+
+RESET search_path;
+-- should fail with wrong number of arguments
+\parameterset work_mem
+\parameterset needs two arguments
+-- should fail because extension is unknown
+\parameterset _pq_.magic_pixie_dust true
+\parameterset cannot be used to change protocol extensions parameters
 -- \gset
 select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_
 \echo :pref01_test01 :pref01_test02 :pref01_test03
@@ -4538,6 +4563,7 @@ invalid command \lo
 	\lo_list
 	\o arg1
 	\p
+	\parameterset abc arg2
 	\password arg1
 	\prompt arg1 arg2
 	\pset arg1 arg2
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
index 129f8533537..1057b83542c 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -59,6 +59,19 @@ SELECT 1 \; SELECT 2 \bind \g
 -- bind error
 SELECT $1, $2 \bind 'foo' \g
 
+-- \parameterset (protocol version 3.1)
+\parameterset work_mem 1234
+SHOW work_mem;
+\parameterset work_mem 'abc'
+RESET work_mem;
+\parameterset search_path '"some quoted string",and_one_without,""'
+SHOW search_path;
+RESET search_path;
+-- should fail with wrong number of arguments
+\parameterset work_mem
+-- should fail because extension is unknown
+\parameterset _pq_.magic_pixie_dust true
+
 -- \gset
 
 select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_
@@ -1020,6 +1033,7 @@ select \if false \\ (bogus \else \\ 42 \endif \\ forty_two;
 	\lo_list
 	\o arg1
 	\p
+	\parameterset abc arg2
 	\password arg1
 	\prompt arg1 arg2
 	\pset arg1 arg2
-- 
2.34.1