psql: Add command to use extended query protocol

This adds a new psql command \bind that sets query parameters and
causes the next query to be sent using the extended query protocol.
Example:

    SELECT $1, $2 \bind 'foo' 'bar' \g

This may be useful for psql scripting, but one of the main purposes is
also to be able to test various aspects of the extended query protocol
from psql and to write tests more easily.

Reviewed-by: Corey Huinker <corey.huinker@gmail.com>
Discussion: https://www.postgresql.org/message-id/flat/e8dd1cd5-0e04-3598-0518-a605159fe314@enterprisedb.com

7 files changed, 101 insertions, 1 deletions

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index ab613dd49e0..3b06169ba0d 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -63,6 +63,7 @@ static backslashResult exec_command(const char *cmd,
 									PQExpBuffer query_buf,
 									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_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,
@@ -308,6 +309,8 @@ exec_command(const char *cmd,
 
 	if (strcmp(cmd, "a") == 0)
 		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, "C") == 0)
 		status = exec_command_C(scan_state, active_branch);
 	else if (strcmp(cmd, "c") == 0 || strcmp(cmd, "connect") == 0)
@@ -454,6 +457,40 @@ exec_command_a(PsqlScanState scan_state, bool active_branch)
 }
 
 /*
+ * \bind -- set query parameters
+ */
+static backslashResult
+exec_command_bind(PsqlScanState scan_state, bool active_branch)
+{
+	backslashResult status = PSQL_CMD_SKIP_LINE;
+
+	if (active_branch)
+	{
+		char	   *opt;
+		int			nparams = 0;
+		int			nalloc = 0;
+
+		pset.bind_params = NULL;
+
+		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] = pg_strdup(opt);
+		}
+
+		pset.bind_nparams = nparams;
+		pset.bind_flag = true;
+	}
+
+	return status;
+}
+
+/*
  * \C -- override table title (formerly change HTML caption)
  */
 static backslashResult
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 864f195992f..b989d792aa7 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1220,6 +1220,16 @@ sendquery_cleanup:
 		pset.gsavepopt = NULL;
 	}
 
+	/* clean up after \bind */
+	if (pset.bind_flag)
+	{
+		for (i = 0; i < pset.bind_nparams; i++)
+			free(pset.bind_params[i]);
+		free(pset.bind_params);
+		pset.bind_params = NULL;
+		pset.bind_flag = false;
+	}
+
 	/* reset \gset trigger */
 	if (pset.gset_prefix)
 	{
@@ -1397,7 +1407,10 @@ ExecQueryAndProcessResults(const char *query,
 	if (timing)
 		INSTR_TIME_SET_CURRENT(before);
 
-	success = PQsendQuery(pset.db, query);
+	if (pset.bind_flag)
+		success = PQsendQueryParams(pset.db, query, pset.bind_nparams, NULL, (const char * const *) pset.bind_params, NULL, NULL, 0);
+	else
+		success = PQsendQuery(pset.db, query);
 
 	if (!success)
 	{
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index f8ce1a07060..b4e0ec2687f 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -189,6 +189,7 @@ slashUsage(unsigned short int pager)
 	initPQExpBuffer(&buf);
 
 	HELP0("General\n");
+	HELP0("  \\bind [PARAM]...       set query parameters\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");
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 2399cffa3fb..3fce71b85fe 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -96,6 +96,9 @@ typedef struct _psqlSettings
 	char	   *gset_prefix;	/* one-shot prefix argument for \gset */
 	bool		gdesc_flag;		/* one-shot request to describe query result */
 	bool		gexec_flag;		/* one-shot request to execute query result */
+	bool		bind_flag;		/* one-shot request to use extended query protocol */
+	int			bind_nparams;	/* number of parameters */
+	char	  **bind_params;	/* parameters for extended query protocol 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 7b73886ce19..a0e26bc2958 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -1680,6 +1680,7 @@ psql_completion(const char *text, int start, int end)
 	/* psql's backslash commands. */
 	static const char *const backslash_commands[] = {
 		"\\a",
+		"\\bind",
 		"\\connect", "\\conninfo", "\\C", "\\cd", "\\copy",
 		"\\copyright", "\\crosstabview",
 		"\\d", "\\da", "\\dA", "\\dAc", "\\dAf", "\\dAo", "\\dAp",
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index a7f5700edc1..5bdae290dce 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -98,6 +98,37 @@ two | 2
    1 |   2
 (1 row)
 
+-- \bind (extended query protocol)
+SELECT 1 \bind \g
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT $1 \bind 'foo' \g
+ ?column? 
+----------
+ foo
+(1 row)
+
+SELECT $1, $2 \bind 'foo' 'bar' \g
+ ?column? | ?column? 
+----------+----------
+ foo      | bar
+(1 row)
+
+-- errors
+-- parse error
+SELECT foo \bind \g
+ERROR:  column "foo" does not exist
+LINE 1: SELECT foo 
+               ^
+-- tcop error
+SELECT 1 \; SELECT 2 \bind \g
+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
 -- \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 1149c6a839e..8732017e51e 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -45,6 +45,20 @@ 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
 
+-- \bind (extended query protocol)
+
+SELECT 1 \bind \g
+SELECT $1 \bind 'foo' \g
+SELECT $1, $2 \bind 'foo' 'bar' \g
+
+-- errors
+-- parse error
+SELECT foo \bind \g
+-- tcop error
+SELECT 1 \; SELECT 2 \bind \g
+-- bind error
+SELECT $1, $2 \bind 'foo' \g
+
 -- \gset
 
 select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_