From bae717f8b2710a6e29f0acf5b52de5f5501b8ccc Mon Sep 17 00:00:00 2001
From: Yugo Nagata <nagata@sraoss.co.jp>
Date: Tue, 12 Aug 2025 01:08:24 +0900
Subject: [PATCH 1/2] pgbench: Allow variables to be used as an SQL literal or
 identifier

Now, the form :'var' and :"var" can be used in addition to :var,
but only in the SQL command. They do not work in arguments of
meta-commands for now.
---
 src/bin/pgbench/pgbench.c | 116 +++++++++++++++++++++++++++++++++-----
 1 file changed, 103 insertions(+), 13 deletions(-)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 125f3c7bbbe..14d1261ff8d 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -1913,10 +1913,32 @@ putVariableInt(Variables *variables, const char *context, char *name,
  * Otherwise, return NULL.
  */
 static char *
-parseVariable(const char *sql, int *eaten)
+parseVariable(const char *sql, int *eaten, PsqlScanQuoteType *quote)
 {
 	int			i = 1;			/* starting at 1 skips the colon */
 	char	   *name;
+	int			start;
+	int			len;
+
+
+	if (sql[i] == '\'')
+	{
+		i++;
+		*quote = PQUOTE_SQL_LITERAL;
+		start = 2;
+	}
+	else if (sql[i] == '"')
+	{
+		i++;
+		*quote = PQUOTE_SQL_IDENT;
+		start = 2;
+	}
+	else
+	{
+		*quote = PQUOTE_PLAIN;
+		start = 1;
+	}
+
 
 	/* keep this logic in sync with valid_variable_name() */
 	if (IS_HIGHBIT_SET(sql[i]) ||
@@ -1931,18 +1953,82 @@ parseVariable(const char *sql, int *eaten)
 				  "_0123456789", sql[i]) != NULL)
 		i++;
 
-	name = pg_malloc(i);
-	memcpy(name, &sql[1], i - 1);
-	name[i - 1] = '\0';
+	if (*quote == PQUOTE_SQL_LITERAL)
+	{
+		if (sql[i] != '\'')
+			return NULL;
+		i++;
+		len = i - 3;
+	}
+	else if (*quote == PQUOTE_SQL_IDENT)
+	{
+		if (sql[i] != '"')
+			return NULL;
+		i++;
+		len = i - 3;
+	}
+	else
+	{
+		Assert(*quote == PQUOTE_PLAIN);
+		len = i - 1;
+	}
+
+	name = pg_malloc(len + 1);
+	memcpy(name, &sql[start], len);
+	name[len] = '\0';
 
 	*eaten = i;
 	return name;
 }
 
 static char *
-replaceVariable(char **sql, char *param, int len, char *value)
+replaceVariable(char **sql, char *param, int len, char *value, PsqlScanQuoteType quote, PGconn *conn)
 {
-	int			valueln = strlen(value);
+	int			valueln;
+	char		quotechar;
+	char	   *quoted_value;
+
+	switch (quote)
+	{
+		case PQUOTE_PLAIN:
+			quoted_value = 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;
+
+				Assert(conn);
+
+				if (quote == PQUOTE_SQL_LITERAL)
+					escaped_value =
+						PQescapeLiteral(conn, value, strlen(value));
+				else
+					escaped_value =
+						PQescapeIdentifier(conn, value, strlen(value));
+
+				if (escaped_value == NULL)
+				{
+					pg_log_error("escape failed: %s", PQerrorMessage(conn));
+					exit(1);
+				}
+
+				/*
+				 * Rather than complicate the lexer's API with a notion of
+				 * which free() routine to use, just pay the price of an extra
+				 * strdup().
+				 */
+				quoted_value = pg_strdup(escaped_value);
+				PQfreemem(escaped_value);
+				break;
+			}
+	}
+
+	valueln = strlen(quoted_value);
 
 	if (valueln > len)
 	{
@@ -1954,13 +2040,15 @@ replaceVariable(char **sql, char *param, int len, char *value)
 
 	if (valueln != len)
 		memmove(param + valueln, param + len, strlen(param + len) + 1);
-	memcpy(param, value, valueln);
+	memcpy(param, quoted_value, valueln);
+
+	pfree(quoted_value);
 
 	return param + valueln;
 }
 
 static char *
-assignVariables(Variables *variables, char *sql)
+assignVariables(Variables *variables, char *sql, PGconn *conn)
 {
 	char	   *p,
 			   *name,
@@ -1970,8 +2058,9 @@ assignVariables(Variables *variables, char *sql)
 	while ((p = strchr(p, ':')) != NULL)
 	{
 		int			eaten;
+		PsqlScanQuoteType quote;
 
-		name = parseVariable(p, &eaten);
+		name = parseVariable(p, &eaten, &quote);
 		if (name == NULL)
 		{
 			while (*p == ':')
@@ -1989,7 +2078,7 @@ assignVariables(Variables *variables, char *sql)
 			continue;
 		}
 
-		p = replaceVariable(&sql, p, eaten, val);
+		p = replaceVariable(&sql, p, eaten, val, quote, conn);
 	}
 
 	return sql;
@@ -3183,7 +3272,7 @@ sendCommand(CState *st, Command *command)
 		char	   *sql;
 
 		sql = pg_strdup(command->argv[0]);
-		sql = assignVariables(&st->variables, sql);
+		sql = assignVariables(&st->variables, sql, st->con);
 
 		pg_log_debug("client %d sending %s", st->id, sql);
 		r = PQsendQuery(st->con, sql);
@@ -5480,8 +5569,9 @@ parseQuery(Command *cmd)
 		char		var[13];
 		char	   *name;
 		int			eaten;
+		PsqlScanQuoteType quote;
 
-		name = parseVariable(p, &eaten);
+		name = parseVariable(p, &eaten, &quote);
 		if (name == NULL)
 		{
 			while (*p == ':')
@@ -5504,7 +5594,7 @@ parseQuery(Command *cmd)
 		}
 
 		sprintf(var, "$%d", cmd->argc);
-		p = replaceVariable(&sql, p, eaten, var);
+		p = replaceVariable(&sql, p, eaten, var, PQUOTE_PLAIN, NULL);
 
 		cmd->argv[cmd->argc] = name;
 		cmd->argc++;
-- 
2.43.0