From 9d8ef98800bd291de145fb1be41f0868546e02ab Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Fri, 3 Apr 2020 11:45:15 +0900 Subject: [PATCH] Add support for \aset in pgbench This option is similar to \gset, except that it is able to store all results from combined SQL queries into separate variables. If a query returns multiple rows, the last result is stored and if a query returns no rows, nothing is stored. While on it, add a TAP test for \gset to check for a failure when a query returns multiple rows. Author: Fabien Coelho Reviewed-by: Ibrar Ahmed, Michael Paquier Discussion: https://postgr.es/m/alpine.DEB.2.21.1904081914200.2529@lancre --- doc/src/sgml/ref/pgbench.sgml | 22 ++++++-- src/bin/pgbench/pgbench.c | 56 ++++++++++++++------ src/bin/pgbench/t/001_pgbench_with_server.pl | 45 ++++++++++++++++ 3 files changed, 104 insertions(+), 19 deletions(-) diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml index 41b3880c91..58a2aa3bf2 100644 --- a/doc/src/sgml/ref/pgbench.sgml +++ b/doc/src/sgml/ref/pgbench.sgml @@ -1057,18 +1057,29 @@ pgbench options d \gset [prefix] + \aset [prefix] - This command may be used to end SQL queries, taking the place of the + These commands may be used to end SQL queries, taking the place of the terminating semicolon (;). - When this command is 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 prefix if provided. + When the \gset command is 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 prefix + if provided. + + + + When the \aset command is used, all combined SQL queries + (separated by \;) have their columns stored into variables + named after column names, and prefixed with prefix + if provided. If a query returns no row, no assignment is made and the variable + can be tested for existence to detect this. If a query returns more than one + row, the last value is kept. @@ -1077,6 +1088,8 @@ pgbench options d p_two and p_three with integers from the third query. The result of the second query is discarded. + The result of the two last combined queries are stored in variables + four and five. UPDATE pgbench_accounts SET abalance = abalance + :delta @@ -1085,6 +1098,7 @@ UPDATE pgbench_accounts -- compound of two queries SELECT 1 \; SELECT 2 AS two, 3 AS three \gset p_ +SELECT 4 AS four \; SELECT 5 AS five \aset diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c index b8864c6ae5..e99af80167 100644 --- a/src/bin/pgbench/pgbench.c +++ b/src/bin/pgbench/pgbench.c @@ -480,6 +480,7 @@ typedef enum MetaCommand META_SHELL, /* \shell */ META_SLEEP, /* \sleep */ META_GSET, /* \gset */ + META_ASET, /* \aset */ META_IF, /* \if */ META_ELIF, /* \elif */ META_ELSE, /* \else */ @@ -504,14 +505,16 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"}; * not applied. * first_line A short, single-line extract of 'lines', for error reporting. * type SQL_COMMAND or META_COMMAND - * meta The type of meta-command, or META_NONE if command is SQL + * meta The type of meta-command, with META_NONE/GSET/ASET if command + * is SQL. * argc Number of arguments of the command, 0 if not yet processed. * argv Command arguments, the first of which is the command or SQL * string itself. For SQL commands, after post-processing * argv[0] is the same as 'lines' with variables substituted. - * varprefix SQL commands terminated with \gset have this set + * varprefix SQL commands terminated with \gset or \aset have this set * to a non NULL value. If nonempty, it's used to prefix the * variable name that receives the value. + * aset do gset on all possible queries of a combined query (\;). * expr Parsed expression, if needed. * stats Time spent in this command. */ @@ -2489,6 +2492,8 @@ getMetaCommand(const char *cmd) mc = META_ENDIF; else if (pg_strcasecmp(cmd, "gset") == 0) mc = META_GSET; + else if (pg_strcasecmp(cmd, "aset") == 0) + mc = META_ASET; else mc = META_NONE; return mc; @@ -2711,17 +2716,25 @@ sendCommand(CState *st, Command *command) * Process query response from the backend. * * If varprefix is not NULL, it's the variable name prefix where to store - * the results of the *last* command. + * the results of the *last* command (META_GSET) or *all* commands + * (META_ASET). * * Returns true if everything is A-OK, false if any error occurs. */ static bool -readCommandResponse(CState *st, char *varprefix) +readCommandResponse(CState *st, MetaCommand meta, char *varprefix) { PGresult *res; PGresult *next_res; int qrynum = 0; + /* + * varprefix should be set only with \gset or \aset, and SQL commands do + * not need it. + */ + Assert((meta == META_NONE && varprefix == NULL) || + ((meta == META_GSET || meta == META_ASET) && varprefix != NULL)); + res = PQgetResult(st->con); while (res != NULL) @@ -2736,7 +2749,7 @@ readCommandResponse(CState *st, char *varprefix) { case PGRES_COMMAND_OK: /* non-SELECT commands */ case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */ - if (is_last && varprefix != NULL) + if (is_last && meta == META_GSET) { pg_log_error("client %d script %d command %d query %d: expected one row, got %d", st->id, st->use_file, st->command, qrynum, 0); @@ -2745,14 +2758,22 @@ readCommandResponse(CState *st, char *varprefix) break; case PGRES_TUPLES_OK: - if (is_last && varprefix != NULL) + if ((is_last && meta == META_GSET) || meta == META_ASET) { - if (PQntuples(res) != 1) + int ntuples = PQntuples(res); + + if (meta == META_GSET && ntuples != 1) { + /* under \gset, report the error */ pg_log_error("client %d script %d command %d query %d: expected one row, got %d", st->id, st->use_file, st->command, qrynum, PQntuples(res)); goto error; } + else if (meta == META_ASET && ntuples <= 0) + { + /* coldly skip empty result under \aset */ + break; + } /* store results into variables */ for (int fld = 0; fld < PQnfields(res); fld++) @@ -2763,9 +2784,9 @@ readCommandResponse(CState *st, char *varprefix) if (*varprefix != '\0') varname = psprintf("%s%s", varprefix, varname); - /* store result as a string */ - if (!putVariable(st, "gset", varname, - PQgetvalue(res, 0, fld))) + /* store last row result as a string */ + if (!putVariable(st, meta == META_ASET ? "aset" : "gset", varname, + PQgetvalue(res, ntuples - 1, fld))) { /* internal error */ pg_log_error("client %d script %d command %d query %d: error storing into variable %s", @@ -3181,7 +3202,9 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg) return; /* don't have the whole result yet */ /* store or discard the query results */ - if (readCommandResponse(st, sql_script[st->use_file].commands[st->command]->varprefix)) + if (readCommandResponse(st, + sql_script[st->use_file].commands[st->command]->meta, + sql_script[st->use_file].commands[st->command]->varprefix)) st->state = CSTATE_END_COMMAND; else st->state = CSTATE_ABORTED; @@ -4660,7 +4683,7 @@ process_backslash_command(PsqlScanState sstate, const char *source) syntax_error(source, lineno, my_command->first_line, my_command->argv[0], "unexpected argument", NULL, -1); } - else if (my_command->meta == META_GSET) + else if (my_command->meta == META_GSET || my_command->meta == META_ASET) { if (my_command->argc > 2) syntax_error(source, lineno, my_command->first_line, my_command->argv[0], @@ -4804,10 +4827,10 @@ ParseScript(const char *script, const char *desc, int weight) if (command) { /* - * If this is gset, merge into the preceding command. (We - * don't use a command slot in this case). + * If this is gset or aset, merge into the preceding command. + * (We don't use a command slot in this case). */ - if (command->meta == META_GSET) + if (command->meta == META_GSET || command->meta == META_ASET) { Command *cmd; @@ -4830,6 +4853,9 @@ ParseScript(const char *script, const char *desc, int weight) else cmd->varprefix = pg_strdup(command->argv[1]); + /* update the sql command meta */ + cmd->meta = command->meta; + /* cleanup unused command */ free_command(command); diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl index b85a3ac32d..e85728c379 100644 --- a/src/bin/pgbench/t/001_pgbench_with_server.pl +++ b/src/bin/pgbench/t/001_pgbench_with_server.pl @@ -699,6 +699,51 @@ SELECT 0 AS i4, 4 AS i4 \gset -- work on the last SQL command under \; \; \; SELECT 0 AS i5 \; SELECT 5 AS i5 \; \; \gset \set i debug(:i5) +} + }); +# \gset cannot accept more than one row, causing command to fail. +pgbench( + '-t 1', 2, + [ qr{type: .*/001_pgbench_gset_two_rows}, qr{processed: 0/1} ], + [qr{expected one row, got 2\b}], + 'pgbench gset command with two rows', + { + '001_pgbench_gset_two_rows' => q{ +SELECT 5432 AS fail UNION SELECT 5433 ORDER BY 1 \gset +} + }); + +# working \aset +# Valid cases. +pgbench( + '-t 1', 0, + [ qr{type: .*/001_pgbench_aset}, qr{processed: 1/1} ], + [ qr{command=3.: int 8\b}, qr{command=4.: int 7\b} ], + 'pgbench aset command', + { + '001_pgbench_aset' => q{ +-- test aset, which applies to a combined query +\; SELECT 6 AS i6 \; SELECT 7 AS i7 \; \aset +-- unless it returns more than one row, last is kept +SELECT 8 AS i6 UNION SELECT 9 ORDER BY 1 DESC \aset +\set i debug(:i6) +\set i debug(:i7) +} + }); +# Empty result set with \aset, causing command to fail. +pgbench( + '-t 1', 2, + [ qr{type: .*/001_pgbench_aset_empty}, qr{processed: 0/1} ], + [ + qr{undefined variable \"i8\"}, + qr{evaluation of meta-command failed\b} + ], + 'pgbench aset command with empty result', + { + '001_pgbench_aset_empty' => q{ +-- empty result +\; SELECT 5432 AS i8 WHERE FALSE \; \aset +\set i debug(:i8) } });