pgbench-into-1.patch

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index f3afedb..8a7ad3e 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -809,6 +809,29 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
   </para>
 
   <variablelist>
+   <varlistentry id='pgbench-metacommand-into'>
+    <term>
+     <literal>\into <replaceable>var1</> [<replaceable>var2</> ...]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      Stores the first fields of the resulting row from the preceding
+      <command>SELECT</> commands into these variables.
+      The queries must yield exactly one row and the number of provided
+      variables must be less than the total number of columns of the results.
+     </para>
+
+     <para>
+      Example:
+<programlisting>
+SELECT abalance FROM pgbench_accounts WHERE aid=5432;
+\into balance
+</programlisting>
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry id='pgbench-metacommand-set'>
     <term>
      <literal>\set <replaceable>varname</> <replaceable>expression</></literal>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 87fb006..d9827dc 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -307,6 +307,7 @@ typedef struct
 	int			type;			/* command type (SQL_COMMAND or META_COMMAND) */
 	int			argc;			/* number of command words */
 	char	   *argv[MAX_ARGS]; /* command word list */
+	char	   *into[MAX_ARGS]; /* NULL-terminated target variables for \into */
 	PgBenchExpr *expr;			/* parsed expression, if needed */
 	SimpleStats stats;			/* time spent in this command */
 } Command;
@@ -1148,6 +1149,96 @@ getQueryParams(CState *st, const Command *command, const char **params)
 		params[i] = getVariable(st, command->argv[i + 1]);
 }
 
+/* read all responses from backend */
+static bool
+read_response(CState *st, char * into[])
+{
+	PGresult   *res;
+	int			i = 0;
+	bool		at_least_one = false;
+
+	while ((res = PQgetResult(st->con)) != NULL)
+	{
+		at_least_one = true;
+
+		switch (PQresultStatus(res))
+		{
+		case PGRES_COMMAND_OK: /* non-SELECT commands */
+		case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
+			break; /* OK */
+
+		case PGRES_TUPLES_OK:
+			if (into != NULL && into[i] != NULL)
+			{
+				/* store result into variables */
+				int ntuples = PQntuples(res),
+					nfields = PQnfields(res),
+					f = 0;
+
+				if (ntuples != 1)
+				{
+					fprintf(stderr,
+							"client %d state %d expecting one row, got %d\n",
+							st->id, st->state, ntuples);
+					st->ecnt++;
+					PQclear(res);
+					discard_response(st);
+					return false;
+				}
+
+				while (into[i] != NULL && f < nfields)
+				{
+					/* store result as a string */
+					if (!putVariable(st, "into", into[i], PQgetvalue(res, 0, f)))
+					{
+						/* internal error, should it rather abort? */
+						fprintf(stderr,
+								"client %d state %d: error storing into var %s\n",
+								st->id, st->state, into[i]);
+						st->ecnt++;
+						PQclear(res);
+						discard_response(st);
+						return false;
+					}
+
+					i++;
+					f++;
+				}
+			}
+			break;	/* OK */
+
+		default:
+			/* everything else is unexpected, so probably an error */
+			fprintf(stderr, "client %d aborted in state %d: %s",
+					st->id, st->state, PQerrorMessage(st->con));
+			st->ecnt++;
+			PQclear(res);
+			discard_response(st);
+			return false;
+		}
+
+		PQclear(res);
+	}
+
+	if (!at_least_one)
+	{
+		fprintf(stderr, "client %d state %d: no results\n", st->id, st->state);
+		st->ecnt++;
+		return false;
+	}
+
+	if (into != NULL && into[i] != NULL)
+	{
+		fprintf(stderr,
+				"client %d state %d: missing results to fill into variable %s\n",
+				st->id, st->state, into[i]);
+		st->ecnt++;
+		return false;
+	}
+
+	return true;
+}
+
 /* get a value as an int, tell if there is a problem */
 static bool
 coerceToInt(PgBenchValue *pval, int64 *ival)
@@ -1764,7 +1855,6 @@ chooseScript(TState *thread)
 static bool
 doCustom(TState *thread, CState *st, StatsData *agg)
 {
-	PGresult   *res;
 	Command   **commands;
 	bool		trans_needs_throttle = false;
 	instr_time	now;
@@ -1891,22 +1981,11 @@ top:
 		{
 			/*
 			 * Read and discard the query result; note this is not included in
-			 * the statement latency numbers.
+			 * the statement latency numbers (above), thus if reading the
+			 * response fails the transaction is counted nevertheless.
 			 */
-			res = PQgetResult(st->con);
-			switch (PQresultStatus(res))
-			{
-				case PGRES_COMMAND_OK:
-				case PGRES_TUPLES_OK:
-					break;		/* OK */
-				default:
-					fprintf(stderr, "client %d aborted in state %d: %s",
-							st->id, st->state, PQerrorMessage(st->con));
-					PQclear(res);
-					return clientDone(st);
-			}
-			PQclear(res);
-			discard_response(st);
+			if (!read_response(st, commands[st->state]->into))
+				return clientDone(st);
 		}
 
 		if (commands[st->state + 1] == NULL)
@@ -2930,6 +3009,24 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 			syntax_error(source, lineno, my_command->line, my_command->argv[0],
 						 "missing command", NULL, -1);
 	}
+	else if (pg_strcasecmp(my_command->argv[0], "into") == 0)
+	{
+		int i;
+
+		/* at least one variable name must be provided */
+		if (my_command->argc < 2)
+			syntax_error(source, lineno, my_command->line, my_command->argv[0],
+						 "missing variable name", NULL, -1);
+
+		/* check that all variable names are valid */
+		for (i = 1; i < my_command->argc; i++)
+		{
+			if (!isLegalVariableName(my_command->argv[i]))
+				syntax_error(source, lineno, my_command->line,
+							 my_command->argv[0], "invalid variable name",
+							 my_command->argv[i], -1);
+		}
+	}
 	else
 	{
 		syntax_error(source, lineno, my_command->line, my_command->argv[0],
@@ -3010,7 +3107,40 @@ ParseScript(const char *script, const char *desc, int weight)
 		if (sr == PSCAN_BACKSLASH)
 		{
 			command = process_backslash_command(sstate, desc);
-			if (command)
+
+			/* special case for \into: merge into preceding SQL command */
+			if (command && command->argc >= 1 &&
+				pg_strcasecmp(command->argv[0], "into") == 0)
+			{
+				if (index == 0)
+				{
+					fprintf(stderr,
+							"meta command \\into cannot start a script\n");
+					exit(1);
+				}
+				else if (ps.commands[index-1]->type != SQL_COMMAND)
+				{
+					fprintf(stderr,
+							"meta command \\into must follow a SELECT (%s)\n",
+							ps.commands[index-1]->line);
+					exit(1);
+				}
+				else if (ps.commands[index-1]->into[0] != NULL)
+				{
+					fprintf(stderr,
+							"SQL command %d already has an associated \\into\n",
+							index-1);
+					exit(1);
+				}
+				/* else (command->argc >= 2): already checked by parser */
+
+				/* into is NULL-terminated thanks to pg_malloc0() */
+				memcpy(ps.commands[index-1]->into,
+					   &command->argv[1], sizeof(char *) * (MAX_ARGS-1));
+
+				pg_free(command);
+			}
+			else if (command)
 			{
 				ps.commands[index] = command;
 				index++;