Hello again,
I'd be okay with SELECT 1, 2 \into one two SELECT 3 \into three
Here is a v2 with more or less this approach, although \into does not end the query, but applies to the current or last sql command. A query is still terminated with a ";".
Now it handles things like :
-- standard sql command
SELECT balance FROM bank WHERE id=1;
\into balance
-- compound sql command, three == 3.
SELECT 1, 2 \; SELECT 3 ;
\into three
-- compound query with 2 selects & 3 variables
SELECT i \into one
FROM generate_series(1, 1) AS i \;
SELECT i+1, i+2 \into two three
FROM generate_series(1, 1) AS i ;
I had to add a few lines in psql scanner to count "\;", so the parsing
logic is a little more complicated than before.
-- Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index f3afedb..cca2cc2 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -809,6 +809,30 @@ 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 current or preceding
+ <command>SELECT</> command 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.
+ This meta command does not end the current SQL command.
+ </para>
+
+ <para>
+ Example:
+<programlisting>
+SELECT abalance \into abalance
+ FROM pgbench_accounts WHERE aid=5432;
+</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..667b63c 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -302,11 +302,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 ***intos; /* per-compound command \into variables */
PgBenchExpr *expr; /* parsed expression, if needed */
SimpleStats stats; /* time spent in this command */
} Command;
@@ -1148,6 +1151,107 @@ 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 ** intos[])
+{
+ PGresult *res;
+ int compound = -1;
+
+ while ((res = PQgetResult(st->con)) != NULL)
+ {
+ compound += 1;
+
+ switch (PQresultStatus(res))
+ {
+ case PGRES_COMMAND_OK: /* non-SELECT commands */
+ if (intos[compound] != NULL)
+ {
+ fprintf(stderr,
+ "client %d state %d compound %d: "
+ "cannot apply \\into to non SELECT statement\n",
+ st->id, st->state, compound);
+ st->ecnt++;
+ return false;
+ }
+ case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
+ break; /* OK */
+
+ case PGRES_TUPLES_OK:
+ if (intos[compound] != NULL)
+ {
+ /* store result into variables */
+ int ntuples = PQntuples(res),
+ nfields = PQnfields(res),
+ f = 0;
+
+ if (ntuples != 1)
+ {
+ fprintf(stderr,
+ "client %d state %d compound %d: "
+ "expecting one row, got %d\n",
+ st->id, st->state, compound, ntuples);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ while (intos[compound][f] != NULL && f < nfields)
+ {
+ /* store result as a string */
+ if (!putVariable(st, "into", intos[compound][f],
+ PQgetvalue(res, 0, f)))
+ {
+ /* internal error, should it rather abort? */
+ fprintf(stderr,
+ "client %d state %d compound %d: "
+ "error storing into var %s\n",
+ st->id, st->state, compound, intos[compound][f]);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ f++;
+ }
+
+ if (intos[compound][f] != NULL)
+ {
+ fprintf(stderr,
+ "client %d state %d compound %d: missing results"
+ " to fill into variable %s\n",
+ st->id, st->state, compound, intos[compound][f]);
+ st->ecnt++;
+ return false;
+ }
+ }
+ break; /* OK */
+
+ default:
+ /* everything else is unexpected, so probably an error */
+ fprintf(stderr, "client %d aborted in state %d compound %d: %s",
+ st->id, st->state, compound, PQerrorMessage(st->con));
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ PQclear(res);
+ }
+
+ if (compound == -1)
+ {
+ fprintf(stderr, "client %d state %d: no results\n", st->id, st->state);
+ 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 +1868,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 +1994,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]->intos))
+ return clientDone(st);
}
if (commands[st->state + 1] == NULL)
@@ -2701,22 +2793,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))
@@ -2736,23 +2816,94 @@ 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->intos = 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 space;
+
+ Assert(my_command->type == SQL_COMMAND && len > 0);
+
+ more = skip_sql_comments(more);
+
+ if (more == NULL)
+ return;
+
+ /* add a separator if needed */
+ space = isspace(my_command->lines[len-1]) ? 0 : 1;
+ lmore = strlen(more);
+ my_command->lines = pg_realloc(my_command->lines, len + space + lmore + 1);
+ if (space > 0)
+ my_command->lines[len] = '\n';
+ memcpy(my_command->lines + len + space, more, lmore+1);
+
+ if (compounds > 0)
+ {
+ int nc = my_command->compound + compounds;
+ my_command->intos =
+ pg_realloc(my_command->intos, sizeof(char**) * (nc+1));
+ memset(my_command->intos + my_command->compound + 1, 0,
+ sizeof(char**) * compounds);
+ 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)
{
- my_command->line = pg_malloc(nlpos - p + 1);
+ my_command->line = pg_malloc(nlpos - p + 1 + 3);
memcpy(my_command->line, p, nlpos - p);
- my_command->line[nlpos - p] = '\0';
+ my_command->line[nlpos - p] = '.';
+ my_command->line[nlpos - p + 1] = '.';
+ my_command->line[nlpos - p + 2] = '.';
+ my_command->line[nlpos - p + 3] = '\0';
}
else
my_command->line = pg_strdup(p);
@@ -2760,19 +2911,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;
}
/*
@@ -2930,6 +3079,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], "into") == 0)
+ {
+ /* 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 names", NULL, -1);
+ }
else
{
syntax_error(source, lineno, my_command->line, my_command->argv[0],
@@ -2941,6 +3097,15 @@ process_backslash_command(PsqlScanState sstate, const char *source)
return my_command;
}
+static bool
+ends_with_semicolon(const char * buf)
+{
+ int i = strlen(buf)-1;
+ while (i > 0 && isspace(buf[i]))
+ i--;
+ return i > 0 && buf[i] == ';';
+}
+
/*
* Parse a script (either the contents of a file, or a built-in script)
* and add it to the list of scripts.
@@ -2953,6 +3118,9 @@ ParseScript(const char *script, const char *desc, int weight)
PQExpBufferData line_buf;
int alloc_num;
int index;
+ bool sql_command_in_progress = false;
+ int lineno;
+ int start_offset;
#define COMMANDS_ALLOC_NUM 128
alloc_num = COMMANDS_ALLOC_NUM;
@@ -2976,6 +3144,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);
@@ -2985,31 +3154,27 @@ 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 (sql_command_in_progress)
{
- 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 was interrupted by an \into */
+ append_sql_command(ps.commands[index-1], line_buf.data,
+ sstate->semicolons);
+ sql_command_in_progress = 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;
@@ -3024,6 +3189,75 @@ ParseScript(const char *script, const char *desc, int weight)
}
}
+ if (sr == PSCAN_BACKSLASH)
+ {
+ command = process_backslash_command(sstate, desc);
+
+ if (command)
+ {
+ /* Merge \into into preceeding SQL command ... */
+ if (strncmp(command->argv[0], "into", 4) == 0)
+ {
+ int cindex, i;
+ Command *into_cmd;
+
+ if (index == 0)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\into cannot start a script",
+ NULL, -1);
+
+ into_cmd = ps.commands[index-1];
+
+ if (into_cmd->type != SQL_COMMAND)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\into must follow an SQL command",
+ into_cmd->line, -1);
+
+ /* this \into applies to this sub-command */
+ cindex = into_cmd->compound -
+ (ends_with_semicolon(line_buf.data) ? 1 : 0);
+
+ if (into_cmd->intos[cindex] != NULL)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\into cannot follow an \\into",
+ NULL, -1);
+
+ /* check that all variable names are valid */
+ for (i = 1; i < command->argc; i++)
+ {
+ if (!isLegalVariableName(command->argv[i]))
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\into invalid variable name",
+ command->argv[i], -1);
+ }
+
+ into_cmd->intos[cindex] =
+ pg_malloc0(sizeof(char *) * (command->argc+1));
+
+ memcpy(into_cmd->intos[cindex], command->argv + 1,
+ sizeof(char*) * command->argc);
+
+ /* cleanup unused backslash command */
+ pg_free(command);
+
+ /* "SELECT ... \into ..." follow-up */
+ sql_command_in_progress = *line_buf.data != '\0';
+ }
+ 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;
@@ -3031,6 +3265,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 ab0f822..4d28fdb 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 55067b4..43b5a4d 100644
--- a/src/fe_utils/psqlscan.l
+++ b/src/fe_utils/psqlscan.l
@@ -680,6 +680,9 @@ other .
"\\"[;:] {
/* Force a semicolon or colon into the query buffer */
+ if (yytext[1] == ';')
+ /* count compound commands */
+ cur_state->semicolons++;
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 a52929d..a184fc7 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 */
/*
into.sql
Description: application/sql
into2.sql
Description: application/sql
into-err-1.sql
Description: application/sql
into-err-2.sql
Description: application/sql
into-err-3.sql
Description: application/sql
into-err-4.sql
Description: application/sql
into-err-5.sql
Description: application/sql
into-err-6.sql
Description: application/sql
into-err-7.sql
Description: application/sql
into-err-8.sql
Description: application/sql
into-err-9.sql
Description: application/sql
into-err-10.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
