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 ([email protected])
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers
