1 files changed, 66 insertions, 32 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;
 }