3 files changed, 113 insertions, 40 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index b06ae9779d5..a2f1259c1e2 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -116,19 +116,19 @@ setQFout(const char *fname)
  * If the specified variable exists, return its value as a string (malloc'd
  * and expected to be freed by the caller); else return NULL.
  *
- * If "escape" is true, return the value suitably quoted and escaped,
- * as an identifier or string literal depending on "as_ident".
- * (Failure in escaping should lead to returning NULL.)
+ * If "quote" isn't PQUOTE_PLAIN, then return the value suitably quoted and
+ * escaped for the specified quoting requirement.  (Failure in escaping
+ * should lead to printing an error and returning NULL.)
  *
  * "passthrough" is the pointer previously given to psql_scan_set_passthrough.
  * In psql, passthrough points to a ConditionalStack, which we check to
  * determine whether variable expansion is allowed.
  */
 char *
-psql_get_variable(const char *varname, bool escape, bool as_ident,
+psql_get_variable(const char *varname, PsqlScanQuoteType quote,
 				  void *passthrough)
 {
-	char	   *result;
+	char	   *result = NULL;
 	const char *value;
 
 	/* In an inactive \if branch, suppress all variable substitutions */
@@ -139,40 +139,74 @@ psql_get_variable(const char *varname, bool escape, bool as_ident,
 	if (!value)
 		return NULL;
 
-	if (escape)
+	switch (quote)
 	{
-		char	   *escaped_value;
+		case PQUOTE_PLAIN:
+			result = pg_strdup(value);
+			break;
+		case PQUOTE_SQL_LITERAL:
+		case PQUOTE_SQL_IDENT:
+			{
+				/*
+				 * For these cases, we use libpq's quoting functions, which
+				 * assume the string is in the connection's client encoding.
+				 */
+				char	   *escaped_value;
 
-		if (!pset.db)
-		{
-			psql_error("cannot escape without active connection\n");
-			return NULL;
-		}
+				if (!pset.db)
+				{
+					psql_error("cannot escape without active connection\n");
+					return NULL;
+				}
 
-		if (as_ident)
-			escaped_value =
-				PQescapeIdentifier(pset.db, value, strlen(value));
-		else
-			escaped_value =
-				PQescapeLiteral(pset.db, value, strlen(value));
+				if (quote == PQUOTE_SQL_LITERAL)
+					escaped_value =
+						PQescapeLiteral(pset.db, value, strlen(value));
+				else
+					escaped_value =
+						PQescapeIdentifier(pset.db, value, strlen(value));
 
-		if (escaped_value == NULL)
-		{
-			const char *error = PQerrorMessage(pset.db);
+				if (escaped_value == NULL)
+				{
+					const char *error = PQerrorMessage(pset.db);
 
-			psql_error("%s", error);
-			return NULL;
-		}
+					psql_error("%s", error);
+					return NULL;
+				}
 
-		/*
-		 * Rather than complicate the lexer's API with a notion of which
-		 * free() routine to use, just pay the price of an extra strdup().
-		 */
-		result = pg_strdup(escaped_value);
-		PQfreemem(escaped_value);
+				/*
+				 * Rather than complicate the lexer's API with a notion of
+				 * which free() routine to use, just pay the price of an extra
+				 * strdup().
+				 */
+				result = pg_strdup(escaped_value);
+				PQfreemem(escaped_value);
+				break;
+			}
+		case PQUOTE_SHELL_ARG:
+			{
+				/*
+				 * For this we use appendShellStringNoError, which is
+				 * encoding-agnostic, which is fine since the shell probably
+				 * is too.  In any case, the only special character is "'",
+				 * which is not known to appear in valid multibyte characters.
+				 */
+				PQExpBufferData buf;
+
+				initPQExpBuffer(&buf);
+				if (!appendShellStringNoError(&buf, value))
+				{
+					psql_error("shell command argument contains a newline or carriage return: \"%s\"\n",
+							   value);
+					free(buf.data);
+					return NULL;
+				}
+				result = buf.data;
+				break;
+			}
+
+			/* No default: we want a compiler warning for missing cases */
 	}
-	else
-		result = pg_strdup(value);
 
 	return result;
 }
diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h
index 3d8b8da7fe4..1ceb8ae386c 100644
--- a/src/bin/psql/common.h
+++ b/src/bin/psql/common.h
@@ -12,11 +12,12 @@
 
 #include "libpq-fe.h"
 #include "fe_utils/print.h"
+#include "fe_utils/psqlscan.h"
 
 extern bool openQueryOutputFile(const char *fname, FILE **fout, bool *is_pipe);
 extern bool setQFout(const char *fname);
 
-extern char *psql_get_variable(const char *varname, bool escape, bool as_ident,
+extern char *psql_get_variable(const char *varname, PsqlScanQuoteType quote,
 				  void *passthrough);
 
 extern void psql_error(const char *fmt,...) pg_attribute_printf(1, 2);
diff --git a/src/bin/psql/psqlscanslash.l b/src/bin/psql/psqlscanslash.l
index 319afdc7441..db7a1b9eead 100644
--- a/src/bin/psql/psqlscanslash.l
+++ b/src/bin/psql/psqlscanslash.l
@@ -242,8 +242,7 @@ other			.
 															 yytext + 1,
 															 yyleng - 1);
 						value = cur_state->callbacks->get_variable(varname,
-																   false,
-																   false,
+																   PQUOTE_PLAIN,
 																   cur_state->cb_passthrough);
 						free(varname);
 
@@ -268,14 +267,16 @@ other			.
 				}
 
 :'{variable_char}+'	{
-					psqlscan_escape_variable(cur_state, yytext, yyleng, false);
+					psqlscan_escape_variable(cur_state, yytext, yyleng,
+											 PQUOTE_SQL_LITERAL);
 					*option_quote = ':';
 					unquoted_option_chars = 0;
 				}
 
 
 :\"{variable_char}+\"	{
-					psqlscan_escape_variable(cur_state, yytext, yyleng, true);
+					psqlscan_escape_variable(cur_state, yytext, yyleng,
+											 PQUOTE_SQL_IDENT);
 					*option_quote = ':';
 					unquoted_option_chars = 0;
 				}
@@ -337,9 +338,8 @@ other			.
 
 <xslashbackquote>{
 	/*
-	 * backticked text: copy everything until next backquote, then evaluate.
-	 *
-	 * XXX Possible future behavioral change: substitute for :VARIABLE?
+	 * backticked text: copy everything until next backquote (expanding
+	 * variable references, but doing nought else), then evaluate.
 	 */
 
 "`"				{
@@ -350,6 +350,44 @@ other			.
 					BEGIN(xslasharg);
 				}
 
+:{variable_char}+	{
+					/* Possible psql variable substitution */
+					if (cur_state->callbacks->get_variable == NULL)
+						ECHO;
+					else
+					{
+						char	   *varname;
+						char	   *value;
+
+						varname = psqlscan_extract_substring(cur_state,
+															 yytext + 1,
+															 yyleng - 1);
+						value = cur_state->callbacks->get_variable(varname,
+																   PQUOTE_PLAIN,
+																   cur_state->cb_passthrough);
+						free(varname);
+
+						if (value)
+						{
+							appendPQExpBufferStr(output_buf, value);
+							free(value);
+						}
+						else
+							ECHO;
+					}
+				}
+
+:'{variable_char}+'	{
+					psqlscan_escape_variable(cur_state, yytext, yyleng,
+											 PQUOTE_SHELL_ARG);
+				}
+
+:'{variable_char}*	{
+					/* Throw back everything but the colon */
+					yyless(1);
+					ECHO;
+				}
+
 {other}|\n		{ ECHO; }
 
 }