diff -cprN head/contrib/pgbench/pgbench.c work/contrib/pgbench/pgbench.c *** head/contrib/pgbench/pgbench.c 2009-12-14 09:21:34.822978000 +0900 --- work/contrib/pgbench/pgbench.c 2009-12-15 10:44:04.586068033 +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 *************** getVariable(CState *st, char *name) *** 430,436 **** return NULL; } ! static int putVariable(CState *st, char *name, char *value) { Variable key, --- 431,441 ---- return NULL; } ! /* ! * Put or replace a variable. Returns false on error. ! * The name must consist of alphabets, numbers, and underscore. ! */ ! static bool putVariable(CState *st, char *name, char *value) { Variable key, *************** putVariable(CState *st, char *name, char *** 450,455 **** --- 455,483 ---- if (var == NULL) { Variable *newvars; + int i; + + /* + * Check for the name only when declaring a new variable to avoid + * overhead. + */ + for (i = 0; name[i]; i++) + { + if (!isalnum((unsigned char) name[i]) && name[i] != '_') + { + const char *command = NULL; + + /* Command might not be available during startup. */ + if (sql_files[st->use_file] && + sql_files[st->use_file][st->state] && + sql_files[st->use_file][st->state]->argc > 0) + command = sql_files[st->use_file][st->state]->argv[0]; + + fprintf(stderr, "%s: invalid variable name '%s'\n", + (command ? command : "(global)"), name); + return false; + } + } if (st->variables) newvars = (Variable *) realloc(st->variables, *************** putVariable(CState *st, char *name, char *** 458,464 **** newvars = (Variable *) malloc(sizeof(Variable)); if (newvars == NULL) ! return false; st->variables = newvars; --- 486,492 ---- newvars = (Variable *) malloc(sizeof(Variable)); if (newvars == NULL) ! goto out_of_memory; st->variables = newvars; *************** putVariable(CState *st, char *name, char *** 467,474 **** var->name = NULL; var->value = NULL; ! if ((var->name = strdup(name)) == NULL ! || (var->value = strdup(value)) == NULL) { free(var->name); free(var->value); --- 495,502 ---- var->name = NULL; var->value = NULL; ! if ((var->name = strdup(name)) == NULL || ! (var->value = strdup(value)) == NULL) { free(var->name); free(var->value); *************** putVariable(CState *st, char *name, char *** 492,497 **** --- 520,529 ---- } return true; + + out_of_memory: + fprintf(stderr, "Couldn't allocate memory for variable\n"); + return false; } static char * *************** getQueryParams(CState *st, const Command *** 590,595 **** --- 622,732 ---- 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)) + { + 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)) + 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: *** 878,886 **** #endif snprintf(res, sizeof(res), "%d", getrand(min, max)); ! if (putVariable(st, argv[1], res) == false) { - fprintf(stderr, "%s: out of memory\n", argv[0]); st->ecnt++; return true; } --- 1015,1022 ---- #endif snprintf(res, sizeof(res), "%d", getrand(min, max)); ! if (!putVariable(st, argv[1], res)) { st->ecnt++; return true; } *************** top: *** 948,956 **** } } ! if (putVariable(st, argv[1], res) == false) { - fprintf(stderr, "%s: out of memory\n", argv[0]); st->ecnt++; return true; } --- 1084,1091 ---- } } ! if (!putVariable(st, argv[1], res)) { st->ecnt++; return true; } *************** top: *** 992,998 **** --- 1127,1160 ---- 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; } *************** init(void) *** 1081,1088 **** 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); executeStatement(con, sql); } --- 1243,1250 ---- 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); executeStatement(con, sql); } *************** process_commands(char *buf) *** 1313,1318 **** --- 1475,1496 ---- 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]); *************** main(int argc, char **argv) *** 1722,1732 **** } *p++ = '\0'; ! if (putVariable(&state[0], optarg, p) == false) ! { ! fprintf(stderr, "Couldn't allocate memory for variable\n"); exit(1); - } } break; case 'F': --- 1900,1907 ---- } *p++ = '\0'; ! if (!putVariable(&state[0], optarg, p)) exit(1); } break; case 'F': *************** main(int argc, char **argv) *** 1806,1816 **** state[i].id = i; for (j = 0; j < state[0].nvariables; j++) { ! if (putVariable(&state[i], state[0].variables[j].name, state[0].variables[j].value) == false) ! { ! fprintf(stderr, "Couldn't allocate memory for variable\n"); exit(1); - } } } } --- 1981,1988 ---- state[i].id = i; for (j = 0; j < state[0].nvariables; j++) { ! if (!putVariable(&state[i], state[0].variables[j].name, state[0].variables[j].value)) exit(1); } } } *************** main(int argc, char **argv) *** 1887,1897 **** snprintf(val, sizeof(val), "%d", scale); for (i = 0; i < nclients; i++) { ! if (putVariable(&state[i], "scale", val) == false) ! { ! fprintf(stderr, "Couldn't allocate memory for variable\n"); exit(1); - } } } --- 2059,2066 ---- snprintf(val, sizeof(val), "%d", scale); for (i = 0; i < nclients; i++) { ! if (!putVariable(&state[i], "scale", val)) exit(1); } } 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-15 09:56:00.983310362 +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: