Add \shell and \setshell meta commands to pgbench.

\shell command runs an external shell command.
\setshell also does the same and sets the result to a variable.

original patch by Michael Paquier with some editorialization by Itagaki,
and reviewed by Greg Smith.

Author: ItGacky

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 @@ doCustom(CState *st, instr_time *conn_time)
 
 			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]);

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: