Hi,
I would like to propose patches to extend variable usage in pgbench scripts.
* 0001: Allow variables to be used as an SQL literal or identifier
Currently, variables in pgbench scripts are expanded always without
quotes, so they cannot be used as SQL literals or identifiers.
This patch allows the use of :'var' and :"var" in addition to :var
However, they can used only in SQL commands, not in the arguments of
meta-commands, since exprscan.l (the lexical scanner for pgbench backslash
commands) cannot currently handle quoted values.
Also, we have to use \aset, \gset or -D option to assign a string to a
varialbe,
since pgbench's \set command cannot recognize text values in its arguments.
* 0002: Add syntax for variable exisitence check
Currently, pgbench does not support :{?var} syntax to check whether the
variable
is defined or not. This patch adds support for this syntax in meta-command
arguments.
This is useful for checking if \aset set the result to the variable or not.
However, it cannot be used in SQL statements for now.
Some of the current limitations described above might be relaxed in the future,
but
for now, I would like to ask for initial feedback on this proposal.
Regards,
Yugo Nagata
--
Yugo Nagata <nag...@sraoss.co.jp>
>From a3f262cb96733cc041a7df1a43fe0f81c71f1474 Mon Sep 17 00:00:00 2001
From: Yugo Nagata <nag...@sraoss.co.jp>
Date: Wed, 13 Aug 2025 19:16:06 +0900
Subject: [PATCH 2/2] pgbench: Add syntax for variable exisitence check
Now, :{?var} can be used in meta-command arguments. It can not
be used in SQL statements.
---
src/bin/pgbench/exprscan.l | 4 +++-
src/bin/pgbench/pgbench.c | 27 +++++++++++++++++++++++++--
2 files changed, 28 insertions(+), 3 deletions(-)
diff --git a/src/bin/pgbench/exprscan.l b/src/bin/pgbench/exprscan.l
index 5747af38cb2..fb14f5276e7 100644
--- a/src/bin/pgbench/exprscan.l
+++ b/src/bin/pgbench/exprscan.l
@@ -182,12 +182,14 @@ notnull [Nn][Oo][Tt][Nn][Uu][Ll][Ll]
{else} { return ELSE_KW; }
{end} { return END_KW; }
-:{alnum}+ {
+:{alnum}+ |
+:\{\?{alnum}+\} {
yylval->str = pg_strdup(yytext + 1);
return VARIABLE;
}
{null} { return NULL_CONST; }
+
{true} {
yylval->bval = true;
return BOOLEAN_CONST;
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 14d1261ff8d..ae0b9ccaf83 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -1920,7 +1920,6 @@ parseVariable(const char *sql, int *eaten, PsqlScanQuoteType *quote)
int start;
int len;
-
if (sql[i] == '\'')
{
i++;
@@ -2958,12 +2957,36 @@ evaluateExpr(CState *st, PgBenchExpr *expr, PgBenchValue *retval)
case ENODE_VARIABLE:
{
Variable *var;
+ char *varname = expr->u.variable.varname;
+ bool is_test_var = false;
+
+ if (varname[0] == '{' && varname[1] == '?')
+ {
+ int len;
+
+ is_test_var = true;
+ varname = pg_strdup(varname + 2);
+ len = strlen(varname);
+ varname[len - 1] = '\0';
+ }
- if ((var = lookupVariable(&st->variables, expr->u.variable.varname)) == NULL)
+ if ((var = lookupVariable(&st->variables, varname)) == NULL)
{
pg_log_error("undefined variable \"%s\"", expr->u.variable.varname);
+ if (is_test_var)
+ {
+ free(varname);
+ setBoolValue(retval, false);
+ return true;
+ }
return false;
}
+ else if (is_test_var)
+ {
+ setBoolValue(retval, true);
+ free(varname);
+ return true;
+ }
if (!makeVariableValue(var))
return false;
--
2.43.0
>From bae717f8b2710a6e29f0acf5b52de5f5501b8ccc Mon Sep 17 00:00:00 2001
From: Yugo Nagata <nag...@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, "e);
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, "e);
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
