pgbenchsetshell.patch

diff --git a/contrib/pgbench/pgbench.c b/contrib/pgbench/pgbench.c
index 8a6437f..c0ab24c 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
@@ -992,7 +993,133 @@ top:
 
 			st->listen = 1;
 		}
+		else if (pg_strcasecmp(argv[0], "setshell") == 0)
+		{
+			int	j,
+				retval,
+				retvalglob = 0;
+			char	res[64],
+				resfin[64];
+			FILE	*respipe = NULL;
+			char    commandLoc[SHELL_COMMAND_SIZE];
+
+			/* construction of the command line with all the transmitted arguments */
+                        retval = snprintf(commandLoc,SHELL_COMMAND_SIZE-1,"%s",argv[2]);
+                        if (retval < 0
+				|| retval > SHELL_COMMAND_SIZE-1)
+			{
+				fprintf(stderr, "Error loading parameter by setshell: too many characters\n");
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+			if (argc > 2)
+			{
+				for (j = 3; j < argc; j++)
+				{
+					char *commandLoc2 = strdup(commandLoc);
+					retval = snprintf(commandLoc,SHELL_COMMAND_SIZE-1,"%s %s", commandLoc2, argv[j]);
+					retvalglob += retval;
+					if (retval < 0
+						|| retvalglob > SHELL_COMMAND_SIZE-1)
+					{
+						fprintf(stderr, "Error loading parameter by setshell: too many characters\n");
+						free(commandLoc2);
+						st->ecnt++;
+						return clientDone(st, false);
+					}
+					free(commandLoc2);
+				}
+			}
+			
+			/*
+			 * Data treatment
+			 * prototype: /setshell aid skewerand +additional arguments
+			 */
+			respipe = popen(commandLoc,"r");
+			if (respipe == NULL)
+			{
+				fprintf(stderr, "%s: error launching shell script", argv[0]);
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+
+			if (fgets(res, 64, respipe) == NULL)
+			{
+				fprintf(stderr, "%s: error getting parameter", argv[0]);
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+			
+			retval = pclose(respipe);
+			if (retval == -1)
+			{
+				fprintf(stderr, "%s: error closing shell script", argv[0]);
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+			/* Transform the parameter into an integer */
+			retval = atoi(res);
+			if (retval == 0)
+			{
+				fprintf(stderr, "%s: error input integer", argv[0]);
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+			/* ready to put the variable */
+			snprintf(resfin, sizeof(resfin), "%d", retval);
 
+			if (putVariable(st, argv[1], resfin) == false)
+			{
+				fprintf(stderr, "%s: out of memory\n", argv[0]);
+				st->ecnt++;
+				return clientDone(st, false);
+			}
+#ifdef DEBUG
+			printf("shell parameter name: %s, value: %s", argv[1], resfin);
+#endif
+			st->listen = 1;
+		}
+		else if (pg_strcasecmp(argv[0], "shell") == 0)
+		{
+			int	j,
+				retval,
+				retvalglob;
+			char	commandLoc[SHELL_COMMAND_SIZE];
+
+			retval = snprintf(commandLoc,SHELL_COMMAND_SIZE-1,"%s",argv[1]);
+			if (retval < 0
+				|| retval > SHELL_COMMAND_SIZE-1)
+			{
+				fprintf(stderr, "Error launching shell command: too many characters\n");
+				st->ecnt++;
+				return true;
+			}
+			retvalglob = retval;
+
+			for (j = 2; j < argc; j++)
+			{
+				char *commandLoc2 = strdup(commandLoc);
+				retval = snprintf(commandLoc,SHELL_COMMAND_SIZE-1,"%s %s", commandLoc2, argv[j]);
+				retvalglob += retval;
+				if (retval < 0
+					|| retvalglob > SHELL_COMMAND_SIZE-1)
+				{
+					fprintf(stderr, "Error launching shell command: too many characters\n");
+					free(commandLoc2);
+					st->ecnt++;
+					return true;
+				}
+				free(commandLoc2);
+			}
+			retval = system(commandLoc);
+			if (retval < 0)
+			{
+				fprintf(stderr, "Error launching shell command: command not launched");
+				st->ecnt++;
+				return true;
+			}
+			st->listen = 1;
+		}
 		goto top;
 	}
 
@@ -1313,6 +1440,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]);