From 5091cba8d0ce27b4ec364a9bd2fff3bb1fe1c6e4 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@gmail.com>
Date: Wed, 1 Nov 2023 10:45:12 +0100
Subject: [PATCH] psql: Add support for prepared stmt with extended protocol

Currently, only unnamed prepared statement is supported by psql with the
\bind command so it's not possible to test named statement creation and
execution through extended protocol.

This commit introduces two additional commands: \parse and \bindx
\parse allows to create a prepared statement through extended protocol.
\bindx allows to bind and execute an existing prepared statement through
extended protocol.
---
 doc/src/sgml/ref/psql-ref.sgml     | 57 ++++++++++++++++++++
 src/bin/psql/command.c             | 85 ++++++++++++++++++++++++++++++
 src/bin/psql/common.c              | 19 +++++--
 src/bin/psql/help.c                |  3 ++
 src/bin/psql/settings.h            |  3 ++
 src/bin/psql/tab-complete.c        |  4 +-
 src/test/regress/expected/psql.out | 28 ++++++++++
 src/test/regress/sql/psql.sql      | 15 ++++++
 8 files changed, 209 insertions(+), 5 deletions(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index d94e3cacfc..24ec30e318 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -916,6 +916,35 @@ INSERT INTO tbl1 VALUES ($1, $2) \bind 'first value' 'second value' \g
        </listitem>
       </varlistentry>
 
+      <varlistentry id="app-psql-meta-command-bindx">
+       <term><literal>\bindx</literal> <replaceable class="parameter">prepared_statement_name</replaceable> [ <replaceable class="parameter">parameter</replaceable> ] ... </term>
+
+       <listitem>
+        <para>
+        <literal>\bindx</literal> is equivalent to <literal>\bind</literal>,
+         except that it takes the name of an explicit existing prepared
+         statement as first parameter.
+        </para>
+
+        <para>
+         Example:
+<programlisting>
+PREPARE stmt1 AS INSERT INTO tbl1 VALUES ($1, $2);
+\bindx stmt1 'first value' 'second value' \g
+</programlisting>
+        </para>
+
+        <para>
+         This command causes the extended query protocol (see <xref
+         linkend="protocol-query-concepts"/>) to be used, unlike normal
+         <application>psql</application> operation, which uses the simple
+         query protocol. So this command can be useful to test the extended
+         query protocol from psql.
+        </para>
+
+       </listitem>
+      </varlistentry>
+
       <varlistentry id="app-psql-meta-command-c-lc">
         <term><literal>\c</literal> or <literal>\connect [ -reuse-previous=<replaceable class="parameter">on|off</replaceable> ] [ <replaceable class="parameter">dbname</replaceable> [ <replaceable class="parameter">username</replaceable> ] [ <replaceable class="parameter">host</replaceable> ] [ <replaceable class="parameter">port</replaceable> ] | <replaceable class="parameter">conninfo</replaceable> ]</literal></term>
         <listitem>
@@ -2775,6 +2804,34 @@ lo_import 152801
         </listitem>
       </varlistentry>
 
+      <varlistentry id="app-psql-meta-command-parse">
+        <term><literal>\parse <replaceable class="parameter">prepared_statement_name</replaceable></literal></term>
+        <listitem>
+        <para>
+        Creates a prepared statement for the next query execution.
+        </para>
+
+        <para>
+         Example:
+<programlisting>
+select 1 \parse stmt1 \g
+</programlisting>
+        </para>
+
+        <para>
+         This command causes the extended query protocol (see <xref
+         linkend="protocol-query-concepts"/>) to be used, unlike normal
+         <application>psql</application> operation, which uses the simple
+         query protocol. A <xref linkend="protocol-message-formats-Parse"/>
+         message will be issued by this command so it can be useful to
+         test the extended query protocol from psql. This command affects
+         only the next query executed; all subsequent queries will use the
+         simple query protocol by default.
+        </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 82cc091568..c3d106f6e2 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -64,6 +64,8 @@ static backslashResult exec_command(const char *cmd,
 									PQExpBuffer previous_buf);
 static backslashResult exec_command_a(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_bind(PsqlScanState scan_state, bool active_branch);
+static backslashResult exec_command_bindx(PsqlScanState scan_state, bool active_branch,
+										  const char *cmd);
 static backslashResult exec_command_C(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_connect(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_cd(PsqlScanState scan_state, bool active_branch,
@@ -116,6 +118,8 @@ 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_parse(PsqlScanState scan_state, bool active_branch,
+										  const char* cmd);
 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);
@@ -311,6 +315,8 @@ exec_command(const char *cmd,
 		status = exec_command_a(scan_state, active_branch);
 	else if (strcmp(cmd, "bind") == 0)
 		status = exec_command_bind(scan_state, active_branch);
+	else if (strcmp(cmd, "bindx") == 0)
+		status = exec_command_bindx(scan_state, active_branch, cmd);
 	else if (strcmp(cmd, "C") == 0)
 		status = exec_command_C(scan_state, active_branch);
 	else if (strcmp(cmd, "c") == 0 || strcmp(cmd, "connect") == 0)
@@ -378,6 +384,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, "parse") == 0)
+		status = exec_command_parse(scan_state, active_branch, cmd);
 	else if (strcmp(cmd, "password") == 0)
 		status = exec_command_password(scan_state, active_branch);
 	else if (strcmp(cmd, "prompt") == 0)
@@ -471,6 +479,7 @@ exec_command_bind(PsqlScanState scan_state, bool active_branch)
 		int			nalloc = 0;
 
 		pset.bind_params = NULL;
+		pset.stmtName = NULL;
 
 		while ((opt = psql_scan_slash_option(scan_state, OT_NORMAL, NULL, false)))
 		{
@@ -490,6 +499,50 @@ exec_command_bind(PsqlScanState scan_state, bool active_branch)
 	return status;
 }
 
+/*
+ * \bindx -- set query parameters for an existing prepared statement
+ */
+static backslashResult
+exec_command_bindx(PsqlScanState scan_state, bool active_branch,
+				   const char* cmd)
+{
+	backslashResult status = PSQL_CMD_SKIP_LINE;
+
+	if (active_branch)
+	{
+		char	   *opt;
+		int			nparams = 0;
+		int			nalloc = 0;
+
+		pset.bind_params = NULL;
+
+		opt = psql_scan_slash_option(scan_state, OT_NORMAL, NULL, false);
+		if (!opt)
+		{
+			pg_log_error("\\%s: missing required argument", cmd);
+			status = PSQL_CMD_ERROR;
+			free(opt);
+		} else {
+			pset.stmtName = opt;
+			pset.bind_flag = true;
+
+			while ((opt = psql_scan_slash_option(scan_state, OT_NORMAL, NULL, false)))
+			{
+				nparams++;
+				if (nparams > nalloc)
+				{
+					nalloc = nalloc ? nalloc * 2 : 1;
+					pset.bind_params = pg_realloc_array(pset.bind_params, char *, nalloc);
+				}
+				pset.bind_params[nparams - 1] = opt;
+			}
+			pset.bind_nparams = nparams;
+		}
+	}
+
+	return status;
+}
+
 /*
  * \C -- override table title (formerly change HTML caption)
  */
@@ -2104,6 +2157,38 @@ exec_command_print(PsqlScanState scan_state, bool active_branch,
 	return PSQL_CMD_SKIP_LINE;
 }
 
+/*
+ * \parse -- parse query
+ */
+static backslashResult
+exec_command_parse(PsqlScanState scan_state, bool active_branch,
+				   const char* cmd)
+{
+	backslashResult status = PSQL_CMD_SKIP_LINE;
+
+	if (active_branch)
+	{
+		char	   *opt = psql_scan_slash_option(scan_state,
+												 OT_NORMAL, NULL, false);
+		pset.stmtName = NULL;
+		if (!opt)
+		{
+			pg_log_error("\\%s: missing required argument", cmd);
+			status = PSQL_CMD_ERROR;
+			free(opt);
+		}
+		else
+		{
+			pset.stmtName = opt;
+			pset.parse_flag = true;
+		}
+	}
+	else
+		ignore_slash_options(scan_state);
+
+	return status;
+}
+
 /*
  * \password -- set user password
  */
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index daabf6f12b..c3790ec62c 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1247,10 +1247,20 @@ sendquery_cleanup:
 		for (i = 0; i < pset.bind_nparams; i++)
 			free(pset.bind_params[i]);
 		free(pset.bind_params);
+		free(pset.stmtName);
 		pset.bind_params = NULL;
+		pset.stmtName = NULL;
 		pset.bind_flag = false;
 	}
 
+	/* clean up after \parse */
+	if (pset.parse_flag)
+	{
+		free(pset.stmtName);
+		pset.stmtName = NULL;
+		pset.parse_flag = false;
+	}
+
 	/* reset \gset trigger */
 	if (pset.gset_prefix)
 	{
@@ -1275,7 +1285,6 @@ sendquery_cleanup:
 	return OK;
 }
 
-
 /*
  * DescribeQuery: describe the result columns of a query, without executing it
  *
@@ -1424,7 +1433,7 @@ ExecQueryAndProcessResults(const char *query,
 	bool		return_early = false;
 	instr_time	before,
 				after;
-	PGresult   *result;
+	PGresult   *result = NULL;
 	FILE	   *gfile_fout = NULL;
 	bool		gfile_is_pipe = false;
 
@@ -1433,8 +1442,12 @@ ExecQueryAndProcessResults(const char *query,
 	else
 		INSTR_TIME_SET_ZERO(before);
 
-	if (pset.bind_flag)
+	if (pset.bind_flag && pset.stmtName == NULL)
 		success = PQsendQueryParams(pset.db, query, pset.bind_nparams, NULL, (const char *const *) pset.bind_params, NULL, NULL, 0);
+	else if (pset.bind_flag && pset.stmtName != NULL)
+		success = PQsendQueryPrepared(pset.db, pset.stmtName, pset.bind_nparams, (const char *const *) pset.bind_params, NULL, NULL, 0);
+	else if (pset.parse_flag)
+		success = PQsendPrepare(pset.db, pset.stmtName, query, 0, NULL);
 	else
 		success = PQsendQuery(pset.db, query);
 
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index 3b2d59e2ee..3acf77b775 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -165,6 +165,8 @@ slashUsage(unsigned short int pager)
 
 	HELP0("General\n");
 	HELP0("  \\bind [PARAM]...       set query parameters\n");
+	HELP0("  \\bindx stmt_name [PARAM]...\n"
+		  "                         set query parameters for an existing prepared statement\n");
 	HELP0("  \\copyright             show PostgreSQL usage and distribution terms\n");
 	HELP0("  \\crosstabview [COLUMNS] execute query and display result in crosstab\n");
 	HELP0("  \\errverbose            show most recent error message at maximum verbosity\n");
@@ -312,6 +314,7 @@ slashUsage(unsigned short int pager)
 			  "                         connect to new database (currently no connection)\n");
 	HELP0("  \\conninfo              display information about current connection\n");
 	HELP0("  \\encoding [ENCODING]   show or set client encoding\n");
+	HELP0("  \\parse STMT_NAME       create a prepared statement\n");
 	HELP0("  \\password [USERNAME]   securely change the password for a user\n");
 	HELP0("\n");
 
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 78e9e9692e..373861c820 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -100,6 +100,9 @@ typedef struct _psqlSettings
 								 * protocol */
 	int			bind_nparams;	/* number of parameters */
 	char	  **bind_params;	/* parameters for extended query protocol call */
+	bool		parse_flag;		/* one-shot request to parse query using extended query
+								 * protocol */
+	char	   *stmtName;		/* stmtName for extended query protocol parse call */
 	bool		crosstab_flag;	/* one-shot request to crosstab result */
 	char	   *ctv_args[4];	/* \crosstabview arguments */
 
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index 93742fc6ac..394534e02a 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -1705,7 +1705,7 @@ psql_completion(const char *text, int start, int end)
 	/* psql's backslash commands. */
 	static const char *const backslash_commands[] = {
 		"\\a",
-		"\\bind",
+		"\\bind", "\\bindx",
 		"\\connect", "\\conninfo", "\\C", "\\cd", "\\copy",
 		"\\copyright", "\\crosstabview",
 		"\\d", "\\da", "\\dA", "\\dAc", "\\dAf", "\\dAo", "\\dAp",
@@ -1723,7 +1723,7 @@ psql_completion(const char *text, int start, int end)
 		"\\if", "\\include", "\\include_relative", "\\ir",
 		"\\list", "\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink",
 		"\\out",
-		"\\password", "\\print", "\\prompt", "\\pset",
+		"\\parse", "\\password", "\\print", "\\prompt", "\\pset",
 		"\\qecho", "\\quit",
 		"\\reset",
 		"\\s", "\\set", "\\setenv", "\\sf", "\\sv",
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index c70205b98a..74bf387055 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -98,6 +98,29 @@ two | 2
    1 |   2
 (1 row)
 
+-- \parse (extended query protocol)
+SELECT 1 \parse stmt1 \g
+SELECT $1 \parse stmt2 \g
+SELECT $1, $2 \parse stmt3 \g
+-- \bindx (extended query protocol)
+\bindx stmt1 \g
+ ?column? 
+----------
+        1
+(1 row)
+
+\bindx stmt2 'foo' \g
+ ?column? 
+----------
+ foo
+(1 row)
+
+\bindx stmt3 'foo' 'bar' \g
+ ?column? | ?column? 
+----------+----------
+ foo      | bar
+(1 row)
+
 -- \bind (extended query protocol)
 SELECT 1 \bind \g
  ?column? 
@@ -129,6 +152,11 @@ 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
+-- bindx error
+\bindx stmt3 'baz' \g
+ERROR:  bind message supplies 1 parameters, but prepared statement "stmt3" requires 2
+\bindx stmt4 'baz' \g
+ERROR:  prepared statement "stmt4" does not exist
 -- \gset
 select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_
 \echo :pref01_test01 :pref01_test02 :pref01_test03
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
index 66ff64a160..8070724ec7 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -45,6 +45,18 @@ SELECT 1 as one, 2 as two \g (format=csv csv_fieldsep='\t')
 SELECT 1 as one, 2 as two \gx (title='foo bar')
 \g
 
+-- \parse (extended query protocol)
+
+SELECT 1 \parse stmt1 \g
+SELECT $1 \parse stmt2 \g
+SELECT $1, $2 \parse stmt3 \g
+
+-- \bindx (extended query protocol)
+
+\bindx stmt1 \g
+\bindx stmt2 'foo' \g
+\bindx stmt3 'foo' 'bar' \g
+
 -- \bind (extended query protocol)
 
 SELECT 1 \bind \g
@@ -58,6 +70,9 @@ SELECT foo \bind \g
 SELECT 1 \; SELECT 2 \bind \g
 -- bind error
 SELECT $1, $2 \bind 'foo' \g
+-- bindx error
+\bindx stmt3 'baz' \g
+\bindx stmt4 'baz' \g
 
 -- \gset
 
-- 
2.39.3 (Apple Git-145)

