Hello Tom,

Please look at changing \into to be a SQL-command-ending backslash
command as we previously discussed.

Done.

There are two variants: \gset & \gcset for compound (end SQL query, set variables, but do not end command, so that several settings are allowed in a compound command, a key feature for performance testing).

Personnally, I find the end-of-query semicolon-replacing syntax ugly. However I'm more interested in feature than in elegance on this one, so I'll put up with it.


I think you will find that the implementation is a great deal simpler that way and doesn't require weird hackery on the shared lexer.

I have removed the "hackery", only counting embedded semicolons remains to keep track of compound queries.

Note that the (somehow buggy and indeed not too clean) hackery was not related to the into syntax, but to detecting empty queries which are silently skipped by the server.

If you won't do that, [...]

I think that I have done what you required.

I have documented the fact that now the feature does not work if compound commands contain empty queries, which is a very minor drawback for a pgbench script anyway.

Attached are the patch, a test script for the feature, and various test scripts to trigger error cases.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 3fb29f8..9a0fb12 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -815,6 +815,51 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
   </para>
 
   <variablelist>
+   <varlistentry id='pgbench-metacommand-gset'>
+    <term>
+     <literal>\gcset [<replaceable>prefix</>]</literal> or
+     <literal>\gset [<replaceable>prefix</>]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      These commands may be used to end SQL queries, replacing a semicolon.
+      <literal>\gcset</> replaces an embedded semicolon (<literal>\;</>) within
+      a compound SQL command, and <literal>\gset</> replaces a final
+      (<literal>;</>) semicolon which ends the SQL command. 
+     </para>
+
+     <para>
+      When these commands are used, the preceding SQL query is expected to
+      return one row, the columns of which are stored into variables named after
+      column names, and prefixed with <replaceable>prefix</> if provided.
+     </para>
+
+     <para>
+      The following example puts the final account balance from the first query
+      into variable <replaceable>abalance</>, and fills variables
+      <replaceable>one</>, <replaceable>two</> and <replaceable>p_three</> with
+      integers from a compound query.
+<programlisting>
+UPDATE pgbench_accounts
+  SET abalance = abalance + :delta
+  WHERE aid = :aid
+  RETURNING abalance \gset
+-- compound of two queries
+SELECT 1 AS one, 2 AS two \gcset
+SELECT 3 AS three \gset p_
+</programlisting>
+     </para>
+
+     <note>
+      <para>
+        <literal>\gcset</> and <literal>\gset</> commands do not work when
+        empty SQL queries appear within a compound SQL command.
+      </para>
+     </note>
+    </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 f6cb5d4..d57eb31 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -375,11 +375,14 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
 
 typedef struct
 {
-	char	   *line;			/* text of command line */
+	char	   *line;			/* first line for short display */
+	char	   *lines;			/* full multi-line text of command */
 	int			command_num;	/* unique index of this Command struct */
 	int			type;			/* command type (SQL_COMMAND or META_COMMAND) */
 	int			argc;			/* number of command words */
 	char	   *argv[MAX_ARGS]; /* command word list */
+	int			compound;       /* last compound command (number of \;) */
+	char	  **gset;           /* per-compound command prefix */
 	PgBenchExpr *expr;			/* parsed expression, if needed */
 	SimpleStats stats;			/* time spent in this command */
 } Command;
@@ -1225,6 +1228,105 @@ 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 **gset)
+{
+	PGresult   *res;
+	int			compound = -1;
+
+	while ((res = PQgetResult(st->con)) != NULL)
+	{
+		compound += 1;
+
+		switch (PQresultStatus(res))
+		{
+		case PGRES_COMMAND_OK: /* non-SELECT commands */
+		case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
+			if (gset[compound] != NULL)
+			{
+				fprintf(stderr,
+						"client %d file %d command %d compound %d: "
+						"\\gset expects a row\n",
+						st->id, st->use_file, st->command, compound);
+				st->ecnt++;
+				return false;
+			}
+			break; /* OK */
+
+		case PGRES_TUPLES_OK:
+			if (gset[compound] != NULL)
+			{
+				/* store result into variables */
+				int ntuples = PQntuples(res),
+					nfields = PQnfields(res),
+					f;
+
+				if (ntuples != 1)
+				{
+					fprintf(stderr,
+							"client %d file %d command %d compound %d: "
+							"expecting one row, got %d\n",
+							st->id, st->use_file, st->command, compound, ntuples);
+					st->ecnt++;
+					PQclear(res);
+					discard_response(st);
+					return false;
+				}
+
+				for (f = 0; f < nfields ; f++)
+				{
+					char *varname = PQfname(res, f);
+					if (*gset[compound] != '\0')
+						varname = psprintf("%s%s", gset[compound], varname);
+
+					/* store result as a string */
+					if (!putVariable(st, "gset", varname,
+									 PQgetvalue(res, 0, f)))
+					{
+						/* internal error, should it rather abort? */
+						fprintf(stderr,
+								"client %d file %d command %d compound %d: "
+								"error storing into var %s\n",
+								st->id, st->use_file, st->command, compound,
+								varname);
+						st->ecnt++;
+						PQclear(res);
+						discard_response(st);
+						return false;
+					}
+
+					if (*gset[compound] != '\0')
+						free(varname);
+				}
+			}
+			break;	/* OK */
+
+		default:
+			/* everything else is unexpected, so probably an error */
+			fprintf(stderr,
+					"client %d file %d aborted in command %d compound %d: %s",
+					st->id, st->use_file, st->command, compound,
+					PQerrorMessage(st->con));
+			st->ecnt++;
+			PQclear(res);
+			discard_response(st);
+			return false;
+		}
+
+		PQclear(res);
+	}
+
+	if (compound == -1)
+	{
+		fprintf(stderr, "client %d command %d: no results\n", st->id, st->command);
+		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)
@@ -1957,7 +2059,6 @@ evaluateSleep(CState *st, int argc, char **argv, int *usecs)
 static void
 doCustom(TState *thread, CState *st, StatsData *agg)
 {
-	PGresult   *res;
 	Command    *command;
 	instr_time	now;
 	bool		end_tx_processed = false;
@@ -2295,26 +2396,12 @@ doCustom(TState *thread, CState *st, StatsData *agg)
 				if (PQisBusy(st->con))
 					return;		/* don't have the whole result yet */
 
-				/*
-				 * Read and discard the query result;
-				 */
-				res = PQgetResult(st->con);
-				switch (PQresultStatus(res))
-				{
-					case PGRES_COMMAND_OK:
-					case PGRES_TUPLES_OK:
-					case PGRES_EMPTY_QUERY:
-						/* OK */
-						PQclear(res);
-						discard_response(st);
-						st->state = CSTATE_END_COMMAND;
-						break;
-					default:
-						commandFailed(st, PQerrorMessage(st->con));
-						PQclear(res);
-						st->state = CSTATE_ABORTED;
-						break;
-				}
+				/* read and discard the query results */
+				if (read_response(st, command->gset))
+					st->state = CSTATE_END_COMMAND;
+				else
+					st->state = CSTATE_ABORTED;
+
 				break;
 
 				/*
@@ -2944,22 +3031,10 @@ syntax_error(const char *source, int lineno,
 	exit(1);
 }
 
-/*
- * Parse a SQL command; return a Command struct, or NULL if it's a comment
- *
- * On entry, psqlscan.l has collected the command into "buf", so we don't
- * really need to do much here except check for comment and set up a
- * Command struct.
- */
-static Command *
-process_sql_command(PQExpBuffer buf, const char *source)
+static char *
+skip_sql_comments(char *p)
 {
-	Command    *my_command;
-	char	   *p;
-	char	   *nlpos;
-
 	/* Skip any leading whitespace, as well as "--" style comments */
-	p = buf->data;
 	for (;;)
 	{
 		if (isspace((unsigned char) *p))
@@ -2979,17 +3054,79 @@ process_sql_command(PQExpBuffer buf, const char *source)
 	if (*p == '\0')
 		return NULL;
 
+	return p;
+}
+
+/*
+ * Parse a SQL command; return a Command struct, or NULL if it's a comment
+ *
+ * On entry, psqlscan.l has collected the command into "buf", so we don't
+ * really need to do much here except check for comment and set up a
+ * Command struct.
+ */
+static Command *
+create_sql_command(PQExpBuffer buf, const char *source, int compounds)
+{
+	Command    *my_command;
+	char	   *p = skip_sql_comments(buf->data);
+
+	if (p == NULL)
+		return NULL;
+
 	/* Allocate and initialize Command structure */
 	my_command = (Command *) pg_malloc0(sizeof(Command));
 	my_command->command_num = num_commands++;
 	my_command->type = SQL_COMMAND;
 	my_command->argc = 0;
+	my_command->compound = compounds;
+	my_command->gset = pg_malloc0(sizeof(char *) * (compounds+1));
 	initSimpleStats(&my_command->stats);
 
+	my_command->lines = pg_strdup(p);
+
+	return my_command;
+}
+
+static void
+append_sql_command(Command *my_command, char *more, int compounds)
+{
+	size_t	lmore;
+	size_t	len = strlen(my_command->lines);
+	int		nc;
+
+	Assert(my_command->type == SQL_COMMAND && len > 0);
+
+	more = skip_sql_comments(more);
+
+	if (more == NULL)
+		return;
+
+	lmore = strlen(more);
+	my_command->lines = pg_realloc(my_command->lines, len + lmore + 2);
+	my_command->lines[len] = ';';
+	memcpy(my_command->lines + len + 1, more, lmore + 1);
+
+	nc = my_command->compound + 1 + compounds;
+	my_command->gset =
+		pg_realloc(my_command->gset, sizeof(char *) * (nc+1));
+	memset(my_command->gset + my_command->compound + 1, 0,
+		   sizeof(char *) * (compounds + 1));
+	my_command->compound = nc;
+}
+
+static void
+postprocess_sql_command(Command *my_command)
+{
+	char	   *nlpos;
+	char	   *p;
+
+	Assert(my_command->type == SQL_COMMAND);
+
 	/*
 	 * If SQL command is multi-line, we only want to save the first line as
 	 * the "line" label.
 	 */
+	p = my_command->lines;
 	nlpos = strchr(p, '\n');
 	if (nlpos)
 	{
@@ -3003,19 +3140,17 @@ process_sql_command(PQExpBuffer buf, const char *source)
 	switch (querymode)
 	{
 		case QUERY_SIMPLE:
-			my_command->argv[0] = pg_strdup(p);
+			my_command->argv[0] = my_command->lines;
 			my_command->argc++;
 			break;
 		case QUERY_EXTENDED:
 		case QUERY_PREPARED:
-			if (!parseQuery(my_command, p))
+			if (!parseQuery(my_command, my_command->lines))
 				exit(1);
 			break;
 		default:
 			exit(1);
 	}
-
-	return my_command;
 }
 
 /*
@@ -3173,6 +3308,13 @@ 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], "gset") == 0 ||
+			 pg_strcasecmp(my_command->argv[0], "gcset") == 0)
+	{
+		if (my_command->argc > 2)
+			syntax_error(source, lineno, my_command->line, my_command->argv[0],
+						 "at most on argument expected", NULL, -1);
+	}
 	else
 	{
 		syntax_error(source, lineno, my_command->line, my_command->argv[0],
@@ -3196,6 +3338,9 @@ ParseScript(const char *script, const char *desc, int weight)
 	PQExpBufferData line_buf;
 	int			alloc_num;
 	int			index;
+	bool		is_compound = false;
+	int			lineno;
+	int			start_offset;
 
 #define COMMANDS_ALLOC_NUM 128
 	alloc_num = COMMANDS_ALLOC_NUM;
@@ -3219,6 +3364,7 @@ ParseScript(const char *script, const char *desc, int weight)
 	 * stdstrings should be true, which is a bit riskier.
 	 */
 	psql_scan_setup(sstate, script, strlen(script), 0, true);
+	start_offset = expr_scanner_offset(sstate) - 1;
 
 	initPQExpBuffer(&line_buf);
 
@@ -3228,31 +3374,28 @@ ParseScript(const char *script, const char *desc, int weight)
 	{
 		PsqlScanResult sr;
 		promptStatus_t prompt;
-		Command    *command;
+		Command    *command = NULL;
 
 		resetPQExpBuffer(&line_buf);
+		lineno = expr_scanner_get_lineno(sstate, start_offset);
+
+		sstate->semicolons = 0;
 
 		sr = psql_scan(sstate, &line_buf, &prompt);
 
-		/* If we collected a SQL command, process that */
-		command = process_sql_command(&line_buf, desc);
-		if (command)
+		if (is_compound)
 		{
-			ps.commands[index] = command;
-			index++;
-
-			if (index >= alloc_num)
-			{
-				alloc_num += COMMANDS_ALLOC_NUM;
-				ps.commands = (Command **)
-					pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
-			}
+			/* a multi-line command ended with \gcset */
+			append_sql_command(ps.commands[index-1], line_buf.data,
+							   sstate->semicolons);
+			is_compound = false;
 		}
-
-		/* If we reached a backslash, process that */
-		if (sr == PSCAN_BACKSLASH)
+		else
 		{
-			command = process_backslash_command(sstate, desc);
+			/* If we collected a new SQL command, process that */
+			command = create_sql_command(&line_buf, desc, sstate->semicolons);
+
+			/* store new command */
 			if (command)
 			{
 				ps.commands[index] = command;
@@ -3267,6 +3410,67 @@ ParseScript(const char *script, const char *desc, int weight)
 			}
 		}
 
+		if (sr == PSCAN_BACKSLASH)
+		{
+			command = process_backslash_command(sstate, desc);
+
+			if (command)
+			{
+				char * bs_cmd = command->argv[0];
+
+				/* merge gset variants into preceeding SQL command */
+				if (pg_strcasecmp(bs_cmd, "gset") == 0 ||
+					pg_strcasecmp(bs_cmd, "gcset") == 0)
+				{
+					int		cindex;
+					Command *sql_cmd;
+
+					is_compound = bs_cmd[1] == 'c';
+
+					if (index == 0)
+						syntax_error(desc, lineno, NULL, NULL,
+									 "\\gset cannot start a script",
+									 NULL, -1);
+
+					sql_cmd = ps.commands[index-1];
+
+					if (sql_cmd->type != SQL_COMMAND)
+						syntax_error(desc, lineno, NULL, NULL,
+									 "\\gset must follow a SQL command",
+									 sql_cmd->line, -1);
+
+					/* this \gset applies to the last sub-command */
+					cindex = sql_cmd->compound;
+
+					if (sql_cmd->gset[cindex] != NULL)
+						syntax_error(desc, lineno, NULL, NULL,
+									 "\\gset cannot follow one another",
+									 NULL, -1);
+
+					/* get variable prefix */
+					if (command->argc <= 1 || command->argv[1][0] == '\0')
+						sql_cmd->gset[cindex] = "";
+					else
+						sql_cmd->gset[cindex] = command->argv[1];
+
+					/* cleanup unused backslash command */
+					pg_free(command);
+				}
+				else /* any other backslash command is a Command */
+				{
+					ps.commands[index] = command;
+					index++;
+
+					if (index >= alloc_num)
+					{
+						alloc_num += COMMANDS_ALLOC_NUM;
+						ps.commands = (Command **)
+							pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
+					}
+				}
+			}
+		}
+
 		/* Done if we reached EOF */
 		if (sr == PSCAN_INCOMPLETE || sr == PSCAN_EOL)
 			break;
@@ -3274,6 +3478,11 @@ ParseScript(const char *script, const char *desc, int weight)
 
 	ps.commands[index] = NULL;
 
+	/* complete SQL command initializations */
+	while (--index >= 0)
+		if (ps.commands[index]->type == SQL_COMMAND)
+			postprocess_sql_command(ps.commands[index]);
+
 	addScript(ps);
 
 	termPQExpBuffer(&line_buf);
diff --git a/src/bin/pgbench/pgbench.h b/src/bin/pgbench/pgbench.h
index 38b3af5..00bf432 100644
--- a/src/bin/pgbench/pgbench.h
+++ b/src/bin/pgbench/pgbench.h
@@ -11,6 +11,7 @@
 #ifndef PGBENCH_H
 #define PGBENCH_H
 
+#include "fe_utils/psqlscan_int.h"
 #include "fe_utils/psqlscan.h"
 
 /*
diff --git a/src/fe_utils/psqlscan.l b/src/fe_utils/psqlscan.l
index 1b29341..d923f70 100644
--- a/src/fe_utils/psqlscan.l
+++ b/src/fe_utils/psqlscan.l
@@ -678,8 +678,15 @@ other			.
 	 * substitution.  We want these before {self}, also.
 	 */
 
-"\\"[;:]		{
-					/* Force a semicolon or colon into the query buffer */
+"\\";			{
+					/* Count compound commands */
+					cur_state->semicolons++;
+					/* Force a semicolon into the query buffer */
+					psqlscan_emit(cur_state, yytext + 1, 1);
+				}
+
+"\\":			{
+					/* Force a colon into the query buffer */
 					psqlscan_emit(cur_state, yytext + 1, 1);
 				}
 
diff --git a/src/include/fe_utils/psqlscan_int.h b/src/include/fe_utils/psqlscan_int.h
index 0fddc7a..c1899df 100644
--- a/src/include/fe_utils/psqlscan_int.h
+++ b/src/include/fe_utils/psqlscan_int.h
@@ -112,6 +112,7 @@ typedef struct PsqlScanStateData
 	int			start_state;	/* yylex's starting/finishing state */
 	int			paren_depth;	/* depth of nesting in parentheses */
 	int			xcdepth;		/* depth of nesting in slash-star comments */
+	int			semicolons;		/* number of embedded (\;) semi-colons */
 	char	   *dolqstart;		/* current $foo$ quote start string */
 
 	/*

Attachment: gset-1.sql
Description: application/sql

Attachment: gset-err-1.sql
Description: application/sql

Attachment: gset-err-2.sql
Description: application/sql

Attachment: gset-err-3.sql
Description: application/sql

Attachment: gset-err-4.sql
Description: application/sql

Attachment: gset-err-5.sql
Description: application/sql

Attachment: gset-err-6.sql
Description: application/sql

-- 
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