diff options
| -rw-r--r-- | contrib/pgbench/pgbench.c | 162 | ||||
| -rw-r--r-- | doc/src/sgml/pgbench.sgml | 52 |
2 files changed, 208 insertions, 6 deletions
diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c index b5520754e4a..470247d0787 100644 --- a/contrib/pgbench/pgbench.c +++ b/contrib/pgbench/pgbench.c @@ -4,7 +4,7 @@ * A simple benchmark program for PostgreSQL * Originally written by Tatsuo Ishii and enhanced by many contributors. * - * $PostgreSQL: pgsql/contrib/pgbench/pgbench.c,v 1.92 2009/12/11 21:50:06 tgl Exp $ + * $PostgreSQL: pgsql/contrib/pgbench/pgbench.c,v 1.93 2009/12/15 07:17:57 itagaki Exp $ * Copyright (c) 2000-2009, PostgreSQL Global Development Group * ALL RIGHTS RESERVED; * @@ -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 @@ -467,8 +468,8 @@ putVariable(CState *st, char *name, char *value) var->name = NULL; var->value = NULL; - if ((var->name = strdup(name)) == NULL - || (var->value = strdup(value)) == NULL) + if ((var->name = strdup(name)) == NULL || + (var->value = strdup(value)) == NULL) { free(var->name); free(var->value); @@ -590,6 +591,114 @@ 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 with whilespace separaters. Arguments starting with + * exactly 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 - 1) + { + 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) + { + if (system(command)) + { + if (!timer_exceeded) + fprintf(stderr, "%s: cannot launch shell command\n", argv[0]); + return false; + } + return true; + } + + /* 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 +1101,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; } @@ -1081,8 +1217,8 @@ init(void) for (i = 0; i < ntellers * scale; i++) { - snprintf(sql, 256, "insert into pgbench_tellers(tid,bid,tbalance) values (%d,%d,0)" - ,i + 1, i / ntellers + 1); + snprintf(sql, 256, "insert into pgbench_tellers(tid,bid,tbalance) values (%d,%d,0)", + i + 1, i / ntellers + 1); executeStatement(con, sql); } @@ -1313,6 +1449,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]); diff --git a/doc/src/sgml/pgbench.sgml b/doc/src/sgml/pgbench.sgml index 0e6ae3f057e..a9cbfe00878 100644 --- a/doc/src/sgml/pgbench.sgml +++ b/doc/src/sgml/pgbench.sgml @@ -1,4 +1,4 @@ -<!-- $PostgreSQL: pgsql/doc/src/sgml/pgbench.sgml,v 1.10 2009/08/03 18:30:55 tgl Exp $ --> +<!-- $PostgreSQL: pgsql/doc/src/sgml/pgbench.sgml,v 1.11 2009/12/15 07:17:57 itagaki Exp $ --> <sect1 id="pgbench"> <title>pgbench</title> @@ -466,6 +466,56 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</> </varlistentry> </variablelist> + <varlistentry> + <term> + <literal>\setshell <replaceable>varname</> <replaceable>command</> [ <replaceable>argument</> ... ]</literal> + </term> + + <listitem> + <para> + Sets variable <replaceable>varname</> to the result of the shell command + <replaceable>command</>. The command must return an integer value + through its standard output. + </para> + + <para> + <replaceable>argument</> can be either a text constant or a + <literal>:</><replaceable>variablename</> reference to a variable of + any types. If you want to use <replaceable>argument</> starting with + colons, you need to add an additional colon at the beginning of + <replaceable>argument</>. + </para> + + <para> + Example: + <programlisting> +\setshell variable_to_be_assigned command literal_argument :variable ::literal_starting_with_colon + </programlisting> + </para> + </listitem> + </varlistentry> + </variablelist> + + <varlistentry> + <term> + <literal>\shell <replaceable>command</> [ <replaceable>argument</> ... ]</literal> + </term> + + <listitem> + <para> + Same as <literal>\setshell</literal>, but the result is ignored. + </para> + + <para> + Example: + <programlisting> +\shell command literal_argument :variable ::literal_starting_with_colon + </programlisting> + </para> + </listitem> + </varlistentry> + </variablelist> + <para> As an example, the full definition of the built-in TPC-B-like transaction is: |