Hi, Yeah the grammar ":variable" is used to set a parameter, the user will feel disorientated if there is no support. The attached patch supports this new case, based on Itagaki-san's last version. I also added a warning to tell the user that pgbench is slowing down when using this feature.
If nobody is against it, I will mark it as ready to commit. -- Regards, Michael Paquier NIPPON TELEGRAPH AND TELEPHONE CORPORATION NTT Open Source Software Center
diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c index 8a6437f..a5b8aeb 100644 --- a/contrib/pgbench/pgbench.c +++ b/contrib/pgbench/pgbench.c @@ -159,6 +159,7 @@ typedef struct } Variable; #define MAX_FILES 128 /* max number of SQL script files allowed */ +#define SHELL_COMMAND_SIZE 256 /* maximum size allowed for shell command */ /* * structures used in custom query mode @@ -590,6 +591,118 @@ getQueryParams(CState *st, const Command *command, const char **params) params[i] = getVariable(st, command->argv[i + 1]); } +/* + * Run a shell command. The result is assigned to the variable if not NULL. + * Return true if succeeded, or false on error. + */ +static bool +runShellCommand(CState *st, char *variable, char **argv, int argc) +{ + char command[SHELL_COMMAND_SIZE]; + int i, + len = 0; + FILE *fp; + char res[64]; + char *endptr; + int retval; + + /* + * Join arguments separated with a whilespace. Arguments starting with + * just one colon are treated as variables: + * name - append a string "name" + * :var - append a variable named 'var'. + * ::name - append a string ":name" + */ + for (i = 0; i < argc; i++) + { + char *arg; + int arglen; + + if (argv[i][0] != ':') + { + arg = argv[i]; /* a string literal */ + } + else if (argv[i][1] == ':') + { + arg = argv[i] + 1; /* a string literal starting with colons */ + } + else if ((arg = getVariable(st, argv[i] + 1)) == NULL) + { + fprintf(stderr, "%s: undefined variable %s\n", argv[0], argv[i]); + return false; + } + + arglen = strlen(arg); + if (len + arglen + (i > 0 ? 1 : 0) >= SHELL_COMMAND_SIZE) + { + fprintf(stderr, "%s: too long shell command\n", argv[0]); + return false; + } + + if (i > 0) + command[len++] = ' '; + memcpy(command + len, arg, arglen); + len += arglen; + } + + command[len] = '\0'; + + /* Fast path for non-assignment case */ + if (variable == NULL) + { + fprintf(stderr, "Warning: slowing down due to %s command overhead\n", argv[0]); + if (system(command)) + { + fprintf(stderr, "%s: cannot launch shell command\n", argv[0]); + return false; + } + return true; + } + + /* Both variable and :variable grammars are supported */ + if (variable[0] == ':') + variable++; + + /* Execute the command with pipe and read the standard output. */ + if ((fp = popen(command, "r")) == NULL) + { + fprintf(stderr, "%s: cannot launch shell command\n", argv[0]); + return false; + } + if (fgets(res, sizeof(res), fp) == NULL) + { + if (!timer_exceeded) + fprintf(stderr, "%s: cannot read the result\n", argv[0]); + return false; + } + if (pclose(fp) < 0) + { + fprintf(stderr, "%s: cannot close shell command\n", argv[0]); + return false; + } + + /* Check whether the result is an integer and assign it to the variable */ + retval = (int) strtol(res, &endptr, 10); + while (*endptr != '\0' && isspace((unsigned char) *endptr)) + endptr++; + if (*res == '\0' || *endptr != '\0') + { + fprintf(stderr, "%s: must return an integer ('%s' returned)\n", argv[0], res); + return false; + } + snprintf(res, sizeof(res), "%d", retval); + if (!putVariable(st, variable, res)) + { + fprintf(stderr, "%s: out of memory\n", argv[0]); + return false; + } + +#ifdef DEBUG + printf("shell parameter name: %s, value: %s\n", argv[1], res); +#endif + return true; +} + #define MAX_PREPARE_NAME 32 static void preparedStatementName(char *buffer, int file, int state) @@ -992,7 +1105,34 @@ top: st->listen = 1; } + else if (pg_strcasecmp(argv[0], "setshell") == 0) + { + bool ret = runShellCommand(st, argv[1], argv + 2, argc - 2); + if (timer_exceeded) /* timeout */ + return clientDone(st, true); + else if (!ret) /* on error */ + { + st->ecnt++; + return true; + } + else /* succeeded */ + st->listen = 1; + } + else if (pg_strcasecmp(argv[0], "shell") == 0) + { + bool ret = runShellCommand(st, NULL, argv + 1, argc - 1); + + if (timer_exceeded) /* timeout */ + return clientDone(st, true); + else if (!ret) /* on error */ + { + st->ecnt++; + return true; + } + else /* succeeded */ + st->listen = 1; + } goto top; } @@ -1313,6 +1453,22 @@ process_commands(char *buf) fprintf(stderr, "%s: extra argument \"%s\" ignored\n", my_commands->argv[0], my_commands->argv[j]); } + else if (pg_strcasecmp(my_commands->argv[0], "setshell") == 0) + { + if (my_commands->argc < 3) + { + fprintf(stderr, "%s: missing argument\n", my_commands->argv[0]); + return NULL; + } + } + else if (pg_strcasecmp(my_commands->argv[0], "shell") == 0) + { + if (my_commands->argc < 1) + { + fprintf(stderr, "%s: missing command\n", my_commands->argv[0]); + return NULL; + } + } else { fprintf(stderr, "Invalid command %s\n", my_commands->argv[0]);
-- Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org) To make changes to your subscription: http://www.postgresql.org/mailpref/pgsql-hackers