Add a \setenv command to psql.

adunstan · adunstan · commit e6d9e2106f0d · 2011-12-04T11:43:38.000-05:00
This can be used to set (or unset) environment variables that will
affect programs called by psql (such as the PAGER), probably most
usefully in a .psqlrc file.

Andrew Dunstan, reviewed by Josh Kupershmidt.
diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
@@ -2241,6 +2241,24 @@ lo_import 152801
       </varlistentry>
 
 
+      <varlistentry>
+        <term><literal>\setenv [ <replaceable class="parameter">name</replaceable> [ <replaceable class="parameter">value</replaceable> ] ]</literal></term>
+
+        <listitem>
+        <para>
+        Sets the environment variable <replaceable
+        class="parameter">name</replaceable> to <replaceable
+        class="parameter">value</replaceable>, or if the
+        <replaceable class="parameter">value</replaceable> is
+        not supplied, unsets the environment variable. Example:
+<programlisting>
+testdb=&gt; <userinput>\setenv PAGER less</userinput>
+testdb=&gt; <userinput>\setenv LESS -imx4F</userinput>
+</programlisting>
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><literal>\sf[+] <replaceable class="parameter">function_description</> </literal></term>
 
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
@@ -1110,6 +1110,51 @@ exec_command(const char *cmd,
 		free(opt0);
 	}
 
+
+	/* \setenv -- set environment command */
+	else if (strcmp(cmd, "setenv") == 0)
+	{
+		char	   *envvar = psql_scan_slash_option(scan_state,
+												  OT_NORMAL, NULL, false);
+		char	   *envval = psql_scan_slash_option(scan_state,
+												  OT_NORMAL, NULL, false);
+
+		if (!envvar)
+		{
+			psql_error("\\%s: missing required argument\n", cmd);
+			success = false;
+		}
+		else if (strchr(envvar,'=') != NULL)
+		{
+			psql_error("\\%s: environment variable name must not contain '='\n",
+					   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 */
+			int         len = strlen(envvar) + strlen(envval) + 1;
+			char	   *newval = pg_malloc(len + 1);
+
+			snprintf(newval, len+1, "%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);
+	}
+
 	/* \sf -- show a function's source code */
 	else if (strcmp(cmd, "sf") == 0 || strcmp(cmd, "sf+") == 0)
 	{
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
@@ -158,7 +158,7 @@ slashUsage(unsigned short int pager)
 {
 	FILE	   *output;
 
-	output = PageOutput(93, pager);
+	output = PageOutput(94, pager);
 
 	/* if you add/remove a line here, change the row count above */
 
@@ -257,6 +257,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, _("  \\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"));
