diff -cprN head/contrib/pgbench/pgbench.c work/contrib/pgbench/pgbench.c *** head/contrib/pgbench/pgbench.c 2009-09-11 09:06:02.464575000 +0900 --- work/contrib/pgbench/pgbench.c 2009-12-10 17:01:11.156147526 +0900 *************** typedef struct *** 159,164 **** --- 159,165 ---- } 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 *************** getQueryParams(CState *st, const Command *** 590,595 **** --- 591,703 ---- 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) + { + if (system(command)) + { + 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) *************** top: *** 992,998 **** --- 1100,1133 ---- 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; } *************** process_commands(char *buf) *** 1313,1318 **** --- 1448,1469 ---- 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 -cprN head/doc/src/sgml/pgbench.sgml work/doc/src/sgml/pgbench.sgml *** head/doc/src/sgml/pgbench.sgml 2009-08-04 09:51:32.798114000 +0900 --- work/doc/src/sgml/pgbench.sgml 2009-12-10 16:14:33.686919000 +0900 *************** pgbench options< *** 466,471 **** --- 466,521 ---- + + + \setshell varname command [ argument ... ] + + + + + Sets variable varname to the result of the shell command + command. The command must return an integer value + through its standard output. + + + + argument can be either a text constant or a + :variablename reference to a variable of + any types. If you want to use argument starting with + colons, you need to add an additional colon at the beginning of + argument. + + + + Example: + + \setshell variable_to_be_assigned command literal_argument :variable ::literal_starting_with_colon + + + + + + + + + \shell command [ argument ... ] + + + + + Same as \setshell, but the result is ignored. + + + + Example: + + \shell command literal_argument :variable ::literal_starting_with_colon + + + + + + As an example, the full definition of the built-in TPC-B-like transaction is: