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

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 41b3880c91f7ce16a70d2c998722b8451d922947..58a2aa3bf200af91511d69a3e2f83278504fe8ee 100644 (file)
--- 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 b8864c6ae53f056200c0fad0da128d8407f482bb..e99af8016752dccbe73a7bf427a93351c112f1e5 100644 (file)
--- 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 b85a3ac32dd3b6b52ccc77874c34ca509b30670b..e85728c3790f5d63c256e5637060231ac7f18166 100644 (file)
--- 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)
 }
    });