From 9d36a386608d7349964e76120e48987e3ec67d04 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Mon, 4 Sep 2017 13:45:20 -0400
Subject: [PATCH] Adjust pgbench to allow non-ASCII characters in variable
 names.

This puts it in sync with psql's notion of what is a valid variable name.
Like psql, we document that "non-Latin letters" are allowed, but actually
any non-ASCII character is accepted.

Fabien Coelho

Discussion: https://postgr.es/m/20170405.094548.1184280384967203518.t-ishii@sraoss.co.jp
---
 doc/src/sgml/ref/pgbench.sgml |  2 ++
 src/bin/pgbench/exprscan.l    |  4 +--
 src/bin/pgbench/pgbench.c     | 47 ++++++++++++++++++++++++++++-------
 3 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 03e1212d50a..f5db8d18d32 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -771,6 +771,8 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
 
   <para>
    There is a simple variable-substitution facility for script files.
+   Variable names must consist of letters (including non-Latin letters),
+   digits, and underscores.
    Variables can be set by the command-line <option>-D</> option,
    explained above, or by the meta commands explained below.
    In addition to any variables preset by <option>-D</> command-line options,
diff --git a/src/bin/pgbench/exprscan.l b/src/bin/pgbench/exprscan.l
index dc1367bbdbc..1862fe40302 100644
--- a/src/bin/pgbench/exprscan.l
+++ b/src/bin/pgbench/exprscan.l
@@ -58,9 +58,9 @@ extern void expr_yyset_column(int column_no, yyscan_t yyscanner);
 %option prefix="expr_yy"
 
 /* Character classes */
-alpha			[a-zA-Z_]
+alpha			[a-zA-Z\200-\377_]
 digit			[0-9]
-alnum			[a-zA-Z0-9_]
+alnum			[A-Za-z\200-\377_0-9]
 /* {space} + {nonspace} + {newline} should cover all characters */
 space			[ \t\r\f\v]
 nonspace		[^ \t\r\f\v\n]
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index ae78c7b1d46..b64e3fbee21 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -1020,19 +1020,38 @@ makeVariableNumeric(Variable *var)
 	return true;
 }
 
-/* check whether the name consists of alphabets, numerals and underscores. */
+/*
+ * Check whether a variable's name is allowed.
+ *
+ * We allow any non-ASCII character, as well as ASCII letters, digits, and
+ * underscore.
+ *
+ * Keep this in sync with the definitions of variable name characters in
+ * "src/fe_utils/psqlscan.l", "src/bin/psql/psqlscanslash.l" and
+ * "src/bin/pgbench/exprscan.l".  Also see parseVariable(), below.
+ *
+ * Note: this static function is copied from "src/bin/psql/variables.c"
+ */
 static bool
-isLegalVariableName(const char *name)
+valid_variable_name(const char *name)
 {
-	int			i;
+	const unsigned char *ptr = (const unsigned char *) name;
+
+	/* Mustn't be zero-length */
+	if (*ptr == '\0')
+		return false;
 
-	for (i = 0; name[i] != '\0'; i++)
+	while (*ptr)
 	{
-		if (!isalnum((unsigned char) name[i]) && name[i] != '_')
+		if (IS_HIGHBIT_SET(*ptr) ||
+			strchr("ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz"
+				   "_0123456789", *ptr) != NULL)
+			ptr++;
+		else
 			return false;
 	}
 
-	return (i > 0);				/* must be non-empty */
+	return true;
 }
 
 /*
@@ -1054,7 +1073,7 @@ lookupCreateVariable(CState *st, const char *context, char *name)
 		 * Check for the name only when declaring a new variable to avoid
 		 * overhead.
 		 */
-		if (!isLegalVariableName(name))
+		if (!valid_variable_name(name))
 		{
 			fprintf(stderr, "%s: invalid variable name: \"%s\"\n",
 					context, name);
@@ -1139,6 +1158,14 @@ putVariableInt(CState *st, const char *context, char *name, int64 value)
 	return putVariableNumber(st, context, name, &val);
 }
 
+/*
+ * Parse a possible variable reference (:varname).
+ *
+ * "sql" points at a colon.  If what follows it looks like a valid
+ * variable name, return a malloc'd string containing the variable name,
+ * and set *eaten to the number of characters consumed.
+ * Otherwise, return NULL.
+ */
 static char *
 parseVariable(const char *sql, int *eaten)
 {
@@ -1148,9 +1175,11 @@ parseVariable(const char *sql, int *eaten)
 	do
 	{
 		i++;
-	} while (isalnum((unsigned char) sql[i]) || sql[i] == '_');
+	} while (IS_HIGHBIT_SET(sql[i]) ||
+			 strchr("ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz"
+					"_0123456789", sql[i]) != NULL);
 	if (i == 1)
-		return NULL;
+		return NULL;			/* no valid variable name chars */
 
 	name = pg_malloc(i);
 	memcpy(name, &sql[1], i - 1);
-- 
2.30.2