pgbenchshell_20091210.patch

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 <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: