pgbenchshell2.5.patch

diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c
index 8a6437f..bb7f468 100644
--- a/contrib/pgbench/pgbench.c
+++ b/contrib/pgbench/pgbench.c
@@ -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
@@ -590,6 +591,137 @@ getQueryParams(CState *st, const Command *command, const char **params)
 		params[i] = getVariable(st, command->argv[i + 1]);
 }
 
+static bool
+process_shellcommand(CState *st, char **argv, int argc, char *commandShell)
+{
+	int		j,
+			retvalglob = 0,
+			retval = 0,
+			arg_min = 0;
+	char	*var;
+
+	/* The minimum number of arguments required is different depending on \setshell or \shell */
+	if (pg_strcasecmp(argv[0], "shell") == 0)
+	{
+		arg_min = 2;
+	}
+	else if (pg_strcasecmp(argv[0], "setshell") == 0)
+	{
+		arg_min = 3;
+	}
+	else
+	{
+		fprintf(stderr, "%s: undefined command type\n", argv[0]);
+		return true;
+	}
+
+	/* construction of the command line with all the transmitted arguments */
+	retval = snprintf(commandShell,SHELL_COMMAND_SIZE-1,"%s",argv[arg_min - 1]);
+	if (retval < 0
+		|| retval > SHELL_COMMAND_SIZE-1)
+	{
+		fprintf(stderr, "%s: error loading shell parameter: too many characters\n",argv[0]);
+		return true;
+	}
+
+	for (j = arg_min; j < argc; j++)
+	{
+		char *commandLoc = strdup(commandShell);
+
+		/* Look at first if the argument is ":"-based or not */
+		if (*argv[j] == ':')
+		{
+			/* Case of a "::"-based argument */
+			if (argv[j][1] == ':')
+			{
+				retval = snprintf(commandShell,SHELL_COMMAND_SIZE-1,"%s %s", commandLoc, argv[j] + 1);
+			}
+			else
+			{
+				/* Case of a ":"-only-based argument */
+				if ((var = getVariable(st, argv[j] + 1)) == NULL)
+				{
+					fprintf(stderr, "%s: undefined variable %s\n", argv[0], argv[j]);
+					return true;
+				}
+				retval = snprintf(commandShell,SHELL_COMMAND_SIZE-1,"%s %s", commandLoc, var);
+			}
+		}
+		else
+		{
+			/* Case when the argument is neither ":" nor "::" based */
+			retval = snprintf(commandShell,SHELL_COMMAND_SIZE-1,"%s %s", commandLoc, argv[j]);
+		}
+		retvalglob += retval;
+		if (retval < 0
+			|| retvalglob > SHELL_COMMAND_SIZE-1)
+		{
+			fprintf(stderr, "%s: error loading shell parameter: too many characters\n", argv[0]);
+			free(commandLoc);
+			return true;
+		}
+		free(commandLoc);
+	}
+	return false;
+}
+
+#define SHELL_ERROR	-1
+#define SHELL_OK	0
+#define SHELL_ALARM	1
+
+static int
+process_shellvariable(CState *st, char **argv ,char *commandShell)
+{
+	int		retval;
+	char	res[64];
+	FILE	*respipe = NULL;
+
+	/*
+	 * Data treatment
+	 * prototype: /setshell aid skewerand +additional arguments
+	 */
+
+	respipe = popen(commandShell,"r");
+	if (respipe == NULL)
+	{
+		fprintf(stderr, "%s: error launching shell script\n", argv[0]);
+		return SHELL_ERROR;
+	}
+	if (fgets(res, sizeof(res), respipe) == NULL)
+	{
+		if (timer_exceeded)
+			return SHELL_ALARM;
+		fprintf(stderr, "%s: error getting parameter\n", argv[0]);
+		return SHELL_ERROR;
+	}
+
+	retval = pclose(respipe);
+	if (retval == -1)
+	{
+		fprintf(stderr, "%s: error closing shell script\n", argv[0]);
+		return SHELL_ERROR;
+	}
+	/* Transform the parameter into an integer */
+	retval = atoi(res);
+	if (retval == 0)
+	{
+		fprintf(stderr, "%s: error input integer\n", argv[0]);
+		return SHELL_ERROR;
+	}
+	/* ready to put the variable */
+	snprintf(res, sizeof(res), "%d", retval);
+
+	if (putVariable(st, argv[1], res) == false)
+	{
+		fprintf(stderr, "%s: out of memory\n", argv[0]);
+		return SHELL_ERROR;
+	}
+#ifdef DEBUG
+	printf("shell parameter name: %s, value: %s\n", argv[1], res);
+#endif
+	return SHELL_OK;
+}
+
 #define MAX_PREPARE_NAME		32
 static void
 preparedStatementName(char *buffer, int file, int state)
@@ -992,7 +1124,53 @@ top:
 
 			st->listen = 1;
 		}
+		else if (pg_strcasecmp(argv[0], "setshell") == 0)
+		{
+			int		retval;
+			bool	status = false;
+			char	commandShell[SHELL_COMMAND_SIZE];
+
+			/* construction of the command line with all the transmitted arguments */
+			status = process_shellcommand(st, argv, argc, commandShell);
+			if (status == true)
+			{
+				st->ecnt++;
+				return true;
+			}
+			/* get as script output the variable and put it */
+			retval = process_shellvariable(st, argv, commandShell);
+			if (retval < 0)
+			{
+				st->ecnt++;
+				return true;
+			}
+			else if (retval > 0)
+				return clientDone(st, true); /* exit correctly, time is out */
+
+			st->listen = 1;
+		}
+		else if (pg_strcasecmp(argv[0], "shell") == 0)
+		{
+			int		retval;
+			bool	status = false;
+			char	commandShell[SHELL_COMMAND_SIZE];
 
+			status = process_shellcommand(st, argv, argc, commandShell);
+			if (status == true)
+			{
+				st->ecnt++;
+				return true;
+			}
+
+			retval = system(commandShell);
+			if (retval < 0)
+			{
+				fprintf(stderr, "Error launching shell command: command not launched\n");
+				st->ecnt++;
+				return true;
+			}
+			st->listen = 1;
+		}
 		goto top;
 	}
 
@@ -1313,6 +1491,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]);
diff --git a/doc/src/sgml/pgbench.sgml b/doc/src/sgml/pgbench.sgml
index f535312..ec5b1d5 100644
--- a/doc/src/sgml/pgbench.sgml
+++ b/doc/src/sgml/pgbench.sgml
@@ -466,6 +466,56 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
    </varlistentry>
   </variablelist>
 
+   <varlistentry>
+    <term>
+     <literal>\setshell <replaceable></> <replaceable>varname</> <replaceable>scriptname</> [ <replaceable>arguments</> ]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      Sets variable <replaceable>varname</> from the output of the call of <replaceable>scriptname</>.
+      It is possible to use as <replaceable>arguments</> variables already set
+      with set or setrandom with the common grammar ":",
+      or to define as <replaceable>arguments</> directly a variable with "::" grammar,
+      this is equivalent to <replaceable>arguments</> beginning with ":".
+      The user is free to set up the arguments depending on the needs for his tests.
+      Possibility to use C, perl or other script types in this feature.
+     </para>
+
+     <para>
+      Example:
+      <programlisting>
+\setshell variable_name script_name :var_previously_defined ::var_undirectly_defined other_script_args
+      </programlisting>
+     </para>
+    </listitem>
+   </varlistentry>
+  </variablelist>
+
+   <varlistentry>
+    <term>
+     <literal>\shell <replaceable></> <replaceable>command</> [ <replaceable>arguments</> ]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      Possibility to launch a shell command directly in
+      pgbench with <replaceable>command</>. As Shell command <replaceable>arguments<\>,
+      it is possible to set an variable with ":" or "::" like in setshell. 
+      The user is free to use as many
+      <replaceable>arguments<\> as he wants depending on the measurement needings.
+     </para>
+
+     <para>
+      Example:
+      <programlisting>
+\shell shell_command shell_arguments :var_previously_defined ::var_undirectly_defined other_script_args
+      </programlisting>
+     </para>
+    </listitem>
+   </varlistentry>
+  </variablelist>
+
   <para>
    As an example, the full definition of the built-in TPC-B-like
    transaction is: