Hi,

I took care of making the 3 cases you mentioned working in the new version
of the patch attached. It worked in my case for both shell and setshell
without any problem.
The code has also been reorganized so as to lighten the process in doCustom.
It looks cleaner on this part.

The only remaining issues are the thread bug and the documentation.
For the bug, I am currently looking into it.
I should take a little bit of time, I don't really know yet from where it
comes exactly...
For the documentation, I'll try to write it a little bit more once the code
issues are solved.

Regards

-- 
Michael Paquier
NIPPON TELEGRAPH AND
TELEPHONE CORPORATION
NTT Open Source Software Center
diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c
index 8a6437f..9c33f7e 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,131 @@ 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;
+}
+
+static bool
+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 true;
+	}
+	if (fgets(res, sizeof(res), respipe) == NULL)
+	{
+		fprintf(stderr, "%s: error getting parameter\n", argv[0]);
+		return true;
+	}
+
+	retval = pclose(respipe);
+	if (retval == -1)
+	{
+		fprintf(stderr, "%s: error closing shell script\n", argv[0]);
+		return true;
+	}
+	/* Transform the parameter into an integer */
+	retval = atoi(res);
+	if (retval == 0)
+	{
+		fprintf(stderr, "%s: error input integer\n", argv[0]);
+		return true;
+	}
+	/* 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 true;
+	}
+#ifdef DEBUG
+	printf("shell parameter name: %s, value: %s\n", argv[1], res);
+#endif
+	return false;
+}
+
 #define MAX_PREPARE_NAME		32
 static void
 preparedStatementName(char *buffer, int file, int state)
@@ -992,7 +1118,50 @@ top:
 
 			st->listen = 1;
 		}
+		else if (pg_strcasecmp(argv[0], "setshell") == 0)
+		{
+			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 */
+			status = process_shellvariable(st, argv, commandShell);
+			if (status == true)
+			{
+				st->ecnt++;
+				return true;
+			}
+
+			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 +1482,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..0782cb1 100644
--- a/doc/src/sgml/pgbench.sgml
+++ b/doc/src/sgml/pgbench.sgml
@@ -466,6 +466,52 @@ 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 set random.
+      Or use <replaceable>arguments</> as arguments for the script call.
+      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 script_arg
+      </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</>. 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
+      </programlisting>
+     </para>
+    </listitem>
+   </varlistentry>
+  </variablelist>
+
   <para>
    As an example, the full definition of the built-in TPC-B-like
    transaction is:
-- 
Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org)
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers

Reply via email to