pgbenchshell_20091215.patch

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 <optional> <replaceable>options<
*** 466,471 ****
--- 466,521 ----
     </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: