From 04059e68ffcd8cf4052ccb6a013f0cf2e0095eb8 Mon Sep 17 00:00:00 2001
From: David Fetter <david@fetter.org>
Date: Wed, 16 Dec 2020 13:17:32 -0800
Subject: [PATCH v1] Implement \gsetenv, similar to \gset
To: hackers
MIME-Version: 1.0
Content-Type: multipart/mixed; boundary="------------2.28.0"

This is a multi-part message in MIME format.
--------------2.28.0
Content-Type: text/plain; charset=UTF-8; format=fixed
Content-Transfer-Encoding: 8bit


In passing, make \setenv without arguments dump the environment in a way
similar to what bare \set and \pset do.

diff --git doc/src/sgml/ref/psql-ref.sgml doc/src/sgml/ref/psql-ref.sgml
index 221a967bfe..97d5a2ed19 100644
--- doc/src/sgml/ref/psql-ref.sgml
+++ doc/src/sgml/ref/psql-ref.sgml
@@ -2297,6 +2297,49 @@ hello 10
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+        <term><literal>\gsetenv [ <replaceable class="parameter">prefix</replaceable> ]</literal></term>
+
+        <listitem>
+        <para>
+         Sends the current query buffer to the server and stores the
+         query's output into environment variables.
+         The query to be executed must return exactly one row.  Each column of
+         the row is stored into a separate variable, named the same as the
+         column.  For example:
+<programlisting>
+=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
+-&gt; <userinput>\gsetenv</userinput>
+=&gt; <userinput>\! echo $var1 $var2</userinput>
+hello 10
+</programlisting>
+        </para>
+        <para>
+         If you specify a <replaceable class="parameter">prefix</replaceable>,
+         that string is prepended to the query's column names to create the
+         variable names to use:
+<programlisting>
+=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
+-&gt; <userinput>\gsetenv result_</userinput>
+=&gt; <userinput>\! echo $result_var1 $result_var2</userinput>
+hello 10
+</programlisting>
+        </para>
+        <para>
+         If a column result is NULL, the corresponding environment variable is unset
+         rather than being set.
+        </para>
+        <para>
+         If the query fails or does not return one row,
+         no variables are changed.
+        </para>
+        <para>
+         If the current query buffer is empty, the most recently sent query
+         is re-executed instead.
+        </para>
+        </listitem>
+      </varlistentry>
+
 
       <varlistentry>
         <term><literal>\gx [ (<replaceable class="parameter">option</replaceable>=<replaceable class="parameter">value</replaceable> [...]) ] [ <replaceable class="parameter">filename</replaceable> ]</literal></term>
diff --git src/bin/psql/command.c src/bin/psql/command.c
index 38b588882d..8b8749aca6 100644
--- src/bin/psql/command.c
+++ src/bin/psql/command.c
@@ -94,7 +94,9 @@ static backslashResult process_command_g_options(char *first_option,
 												 const char *cmd);
 static backslashResult exec_command_gdesc(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_gexec(PsqlScanState scan_state, bool active_branch);
-static backslashResult exec_command_gset(PsqlScanState scan_state, bool active_branch);
+static backslashResult exec_command_gset(PsqlScanState scan_state,
+										 bool active_branch,
+										 GSET_TARGET gset_target);
 static backslashResult exec_command_help(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_html(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_include(PsqlScanState scan_state, bool active_branch,
@@ -346,7 +348,9 @@ exec_command(const char *cmd,
 	else if (strcmp(cmd, "gexec") == 0)
 		status = exec_command_gexec(scan_state, active_branch);
 	else if (strcmp(cmd, "gset") == 0)
-		status = exec_command_gset(scan_state, active_branch);
+		status = exec_command_gset(scan_state, active_branch, GSET_TARGET_PSET);
+	else if (strcmp(cmd, "gsetenv") == 0)
+		status = exec_command_gset(scan_state, active_branch, GSET_TARGET_ENV);
 	else if (strcmp(cmd, "h") == 0 || strcmp(cmd, "help") == 0)
 		status = exec_command_help(scan_state, active_branch);
 	else if (strcmp(cmd, "H") == 0 || strcmp(cmd, "html") == 0)
@@ -1451,10 +1455,11 @@ exec_command_gexec(PsqlScanState scan_state, bool active_branch)
 }
 
 /*
- * \gset [prefix] -- send query and store result into variables
+ * \gset[env] [prefix] -- send query and store result into variables
  */
 static backslashResult
-exec_command_gset(PsqlScanState scan_state, bool active_branch)
+exec_command_gset(PsqlScanState scan_state, bool active_branch,
+				  GSET_TARGET gset_target)
 {
 	backslashResult status = PSQL_CMD_SKIP_LINE;
 
@@ -1463,6 +1468,8 @@ exec_command_gset(PsqlScanState scan_state, bool active_branch)
 		char	   *prefix = psql_scan_slash_option(scan_state,
 													OT_NORMAL, NULL, false);
 
+		pset.gset_target = gset_target;
+
 		if (prefix)
 			pset.gset_prefix = prefix;
 		else
@@ -2275,39 +2282,8 @@ exec_command_setenv(PsqlScanState scan_state, bool active_branch,
 													OT_NORMAL, NULL, false);
 		char	   *envval = psql_scan_slash_option(scan_state,
 													OT_NORMAL, NULL, false);
+		success = SetEnvVariable(cmd, envvar, envval);
 
-		if (!envvar)
-		{
-			pg_log_error("\\%s: missing required argument", cmd);
-			success = false;
-		}
-		else if (strchr(envvar, '=') != NULL)
-		{
-			pg_log_error("\\%s: environment variable name must not contain \"=\"",
-						 cmd);
-			success = false;
-		}
-		else if (!envval)
-		{
-			/* No argument - unset the environment variable */
-			unsetenv(envvar);
-			success = true;
-		}
-		else
-		{
-			/* Set variable to the value of the next argument */
-			char	   *newval;
-
-			newval = psprintf("%s=%s", envvar, envval);
-			putenv(newval);
-			success = true;
-
-			/*
-			 * Do not free newval here, it will screw up the environment if
-			 * you do. See putenv man page for details. That means we leak a
-			 * bit of memory here, but not enough to worry about.
-			 */
-		}
 		free(envvar);
 		free(envval);
 	}
diff --git src/bin/psql/common.c src/bin/psql/common.c
index dfbc22970f..2988613a62 100644
--- src/bin/psql/common.c
+++ src/bin/psql/common.c
@@ -28,6 +28,7 @@
 #include "fe_utils/string_utils.h"
 #include "portability/instr_time.h"
 #include "settings.h"
+#include "variables.h"
 
 static bool DescribeQuery(const char *query, double *elapsed_msec);
 static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
@@ -762,15 +763,24 @@ static bool
 StoreQueryTuple(const PGresult *result)
 {
 	bool		success = true;
+	char		*gset_suffix;
+
+	if (pset.gset_target == GSET_TARGET_PSET)
+		gset_suffix = "";
+	else if (pset.gset_target == GSET_TARGET_ENV)
+		gset_suffix = "env";
+	else
+		pg_log_error("How did you manage to get gset_target to be \"%d\%?!?",
+					 pset.gset_target);
 
 	if (PQntuples(result) < 1)
 	{
-		pg_log_error("no rows returned for \\gset");
+		pg_log_error("no rows returned for \\gset%s", gset_suffix);
 		success = false;
 	}
 	else if (PQntuples(result) > 1)
 	{
-		pg_log_error("more than one row returned for \\gset");
+		pg_log_error("more than one row returned for \\gset%s", gset_suffix);
 		success = false;
 	}
 	else
@@ -786,7 +796,7 @@ StoreQueryTuple(const PGresult *result)
 			/* concatenate prefix and column name */
 			varname = psprintf("%s%s", pset.gset_prefix, colname);
 
-			if (VariableHasHook(pset.vars, varname))
+			if (pset.gset_target == GSET_TARGET_PSET && VariableHasHook(pset.vars, varname))
 			{
 				pg_log_warning("attempt to \\gset into specially treated variable \"%s\" ignored",
 							   varname);
@@ -801,14 +811,31 @@ StoreQueryTuple(const PGresult *result)
 				value = NULL;
 			}
 
-			if (!SetVariable(pset.vars, varname, value))
+			if (pset.gset_target == GSET_TARGET_PSET)
 			{
+				if (!SetVariable(pset.vars, varname, value))
+				{
+					free(varname);
+					success = false;
+					continue;
+				}
+
 				free(varname);
-				success = false;
-				break;
 			}
+			else /* Only GSET_TARGET_ENV is possible now */
+			{
+				const char *cmd;
+
+				cmd = "gsetenv";
+				if (!SetEnvVariable(cmd, varname, value))
+				{
+					free(varname);
+					success = false;
+					continue;
+				}
 
-			free(varname);
+				free(varname);
+			}
 		}
 	}
 
diff --git src/bin/psql/help.c src/bin/psql/help.c
index af829282e6..2c135f4abb 100644
--- src/bin/psql/help.c
+++ src/bin/psql/help.c
@@ -180,6 +180,7 @@ slashUsage(unsigned short int pager)
 	fprintf(output, _("  \\gdesc                 describe result of query, without executing it\n"));
 	fprintf(output, _("  \\gexec                 execute query, then execute each value in its result\n"));
 	fprintf(output, _("  \\gset [PREFIX]         execute query and store results in psql variables\n"));
+	fprintf(output, _("  \\gsetenv [PREFIX]      execute query and store results in environment variables\n"));
 	fprintf(output, _("  \\gx [(OPTIONS)] [FILE] as \\g, but forces expanded output mode\n"));
 	fprintf(output, _("  \\q                     quit psql\n"));
 	fprintf(output, _("  \\watch [SEC]           execute query every SEC seconds\n"));
@@ -309,7 +310,7 @@ slashUsage(unsigned short int pager)
 
 	fprintf(output, _("Operating System\n"));
 	fprintf(output, _("  \\cd [DIR]              change the current working directory\n"));
-	fprintf(output, _("  \\setenv NAME [VALUE]   set or unset environment variable\n"));
+	fprintf(output, _("  \\setenv NAME [VALUE]   set environment variable, or list all if no parameters\n"));
 	fprintf(output, _("  \\timing [on|off]       toggle timing of commands (currently %s)\n"),
 			ON(pset.timing));
 	fprintf(output, _("  \\! [COMMAND]           execute command in shell or start interactive shell\n"));
diff --git src/bin/psql/settings.h src/bin/psql/settings.h
index 9601f6e90c..f83e1324aa 100644
--- src/bin/psql/settings.h
+++ src/bin/psql/settings.h
@@ -62,6 +62,13 @@ typedef enum
 	PSQL_COMP_CASE_LOWER
 } PSQL_COMP_CASE;
 
+typedef enum
+{
+	GSET_TARGET_NONE,
+	GSET_TARGET_PSET,
+	GSET_TARGET_ENV
+} GSET_TARGET;
+
 typedef enum
 {
 	hctl_none = 0,
@@ -93,7 +100,7 @@ typedef struct _psqlSettings
 	char	   *gfname;			/* one-shot file output argument for \g */
 	printQueryOpt *gsavepopt;	/* if not null, saved print format settings */
 
-	char	   *gset_prefix;	/* one-shot prefix argument for \gset */
+	char	   *gset_prefix;	/* one-shot prefix argument for \gset* */
 	bool		gdesc_flag;		/* one-shot request to describe query results */
 	bool		gexec_flag;		/* one-shot request to execute query results */
 	bool		crosstab_flag;	/* one-shot request to crosstab results */
@@ -142,6 +149,7 @@ typedef struct _psqlSettings
 	PSQL_ECHO_HIDDEN echo_hidden;
 	PSQL_ERROR_ROLLBACK on_error_rollback;
 	PSQL_COMP_CASE comp_case;
+	GSET_TARGET gset_target;
 	HistControl histcontrol;
 	const char *prompt1;
 	const char *prompt2;
diff --git src/bin/psql/tab-complete.c src/bin/psql/tab-complete.c
index 3a43c09bf6..924abd150b 100644
--- src/bin/psql/tab-complete.c
+++ src/bin/psql/tab-complete.c
@@ -1504,7 +1504,7 @@ psql_completion(const char *text, int start, int end)
 		"\\e", "\\echo", "\\ef", "\\elif", "\\else", "\\encoding",
 		"\\endif", "\\errverbose", "\\ev",
 		"\\f",
-		"\\g", "\\gdesc", "\\gexec", "\\gset", "\\gx",
+		"\\g", "\\gdesc", "\\gexec", "\\gset", "\\gsetenv", "\\gx",
 		"\\h", "\\help", "\\H",
 		"\\i", "\\if", "\\ir",
 		"\\l", "\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink",
diff --git src/bin/psql/variables.c src/bin/psql/variables.c
index 84349fc753..26961d8c23 100644
--- src/bin/psql/variables.c
+++ src/bin/psql/variables.c
@@ -5,6 +5,8 @@
  *
  * src/bin/psql/variables.c
  */
+#include <unistd.h>
+
 #include "postgres_fe.h"
 
 #include "common.h"
@@ -295,6 +297,67 @@ SetVariable(VariableSpace space, const char *name, const char *value)
 	return true;
 }
 
+/*
+ * Common code for of \{,g}setenv
+ *
+ * Set the environment variable named "name" to value "value",
+ * or delete it if "value" is NULL.
+ */
+bool SetEnvVariable(const char *cmd, const char *name, const char *value)
+{
+	bool	success = true;
+	if (!name)
+	{
+		int cmp;
+
+		cmp = strcmp(cmd, "setenv");
+
+		if (cmp == 0)
+		{
+			int i;
+
+			for (i = 0; environ[i] != NULL; i++)
+			{
+				printf("%-24s\n", environ[i]);
+			}
+			success = true;
+		}
+		else
+		{
+			pg_log_error("\\%s: missing required argument", cmd);
+			success = false;
+		}
+	}
+	else if (strchr(name, '=') != NULL)
+	{
+		pg_log_error("\\%s: environment variable name must not contain \"=\"",
+					 cmd);
+		success = false;
+	}
+	else if (!value)
+	{
+		/* No argument - unset the environment variable */
+		unsetenv(name);
+		success = true;
+	}
+	else
+	{
+		/* Set variable to the value of the next argument */
+		char	   *newval;
+
+		newval = psprintf("%s=%s", name, value);
+		putenv(newval);
+		success = true;
+
+		/*
+		 * Do not free newvalue here, as it may screw up the environment if
+		 * you do. See putenv man page for details. That means we leak a
+		 * bit of memory here, but not enough to worry about.
+		 */
+	}
+	return success;
+}
+
 /*
  * Attach substitute and/or assign hook functions to the named variable.
  * If you need only one hook, pass NULL for the other.
diff --git src/bin/psql/variables.h src/bin/psql/variables.h
index b932472021..3f54648fd2 100644
--- src/bin/psql/variables.h
+++ src/bin/psql/variables.h
@@ -86,6 +86,7 @@ void		PrintVariables(VariableSpace space);
 bool		SetVariable(VariableSpace space, const char *name, const char *value);
 bool		SetVariableBool(VariableSpace space, const char *name);
 bool		DeleteVariable(VariableSpace space, const char *name);
+bool		SetEnvVariable(const char *cmd, const char *name, const char *value);
 
 void		SetVariableHooks(VariableSpace space, const char *name,
 							 VariableSubstituteHook shook,
diff --git src/test/regress/expected/psql.out src/test/regress/expected/psql.out
index 7204fdb0b4..82e8f435b1 100644
--- src/test/regress/expected/psql.out
+++ src/test/regress/expected/psql.out
@@ -152,6 +152,45 @@ more than one row returned for \gset
 select 10 as test01, 20 as test02 from generate_series(1,0) \gset
 no rows returned for \gset
 \unset FETCH_COUNT
+-- \gsetenv
+select 10 as test01, 20 as test02, 'Hello' as test03 \gsetenv pref01_
+\! echo $pref01_test01 $pref01_test02 $pref01_test03
+10 20 Hello
+-- should fail: bad environment variable name
+select 10 as "bad=name"
+\gsetenv
+\gsetenv: environment variable name must not contain "="
+-- multiple backslash commands in one line
+select 1 as x, 2 as y \gsetenv pref01_ \\ \! echo $pref01_x
+1
+select 3 as x, 4 as y \gsetenv pref01_ \\ \g \! echo $pref01_x $pref01_y
+ x | y 
+---+---
+ 3 | 4
+(1 row)
+
+3 4
+-- NULL should unset the variable
+\setenv var2 xyz
+select 1 as var1, NULL as var2, 3 as var3 \gsetenv
+\! echo $var1 $var2 $var3
+1 3
+-- \gsetenv requires just one tuple
+select 10 as test01, 20 as test02 from generate_series(1,3) \gsetenv
+more than one row returned for \gsetenv
+select 10 as test01, 20 as test02 from generate_series(1,0) \gsetenv
+no rows returned for \gsetenv
+-- \gsetenv should work in FETCH_COUNT mode too
+\set FETCH_COUNT 1
+select 1 as x, 2 as y \gsetenv pref01_ \\ \! echo $pref01_x
+1
+select 3 as x, 4 as y \gsetenv pref01_ \! echo $pref01_x $pref01_y
+3 4
+select 10 as test01, 20 as test02 from generate_series(1,3) \gsetenv
+more than one row returned for \gsetenv
+select 10 as test01, 20 as test02 from generate_series(1,0) \gsetenv
+no rows returned for \gsetenv
+\unset FETCH_COUNT
 -- \gdesc
 SELECT
     NULL AS zero,
diff --git src/test/regress/sql/psql.sql src/test/regress/sql/psql.sql
index 537d5332aa..a7e9ac77b3 100644
--- src/test/regress/sql/psql.sql
+++ src/test/regress/sql/psql.sql
@@ -83,6 +83,39 @@ select 10 as test01, 20 as test02 from generate_series(1,0) \gset
 
 \unset FETCH_COUNT
 
+-- \gsetenv
+
+select 10 as test01, 20 as test02, 'Hello' as test03 \gsetenv pref01_
+
+\! echo $pref01_test01 $pref01_test02 $pref01_test03
+
+-- should fail: bad environment variable name
+select 10 as "bad=name"
+\gsetenv
+
+-- multiple backslash commands in one line
+select 1 as x, 2 as y \gsetenv pref01_ \\ \! echo $pref01_x
+select 3 as x, 4 as y \gsetenv pref01_ \\ \g \! echo $pref01_x $pref01_y
+
+-- NULL should unset the variable
+\setenv var2 xyz
+select 1 as var1, NULL as var2, 3 as var3 \gsetenv
+\! echo $var1 $var2 $var3
+
+-- \gsetenv requires just one tuple
+select 10 as test01, 20 as test02 from generate_series(1,3) \gsetenv
+select 10 as test01, 20 as test02 from generate_series(1,0) \gsetenv
+
+-- \gsetenv should work in FETCH_COUNT mode too
+\set FETCH_COUNT 1
+
+select 1 as x, 2 as y \gsetenv pref01_ \\ \! echo $pref01_x
+select 3 as x, 4 as y \gsetenv pref01_ \! echo $pref01_x $pref01_y
+select 10 as test01, 20 as test02 from generate_series(1,3) \gsetenv
+select 10 as test01, 20 as test02 from generate_series(1,0) \gsetenv
+
+\unset FETCH_COUNT
+
 -- \gdesc
 
 SELECT

--------------2.28.0--


