Hi,

The psql improvement in v15 to output multiple result sets does not
behave as one might expect with \g: the output file or program
to pipe into is opened/closed on each result set, overwriting the
previous ones in the case of \g file.

Example:

psql -At <<EOF
-- good (two results output)
select 1\; select 2;

-- bad: ends up with only "2" in the file
select 1\; select 2 \g file

EOF


That problem with \g is due to PrintQueryTuples() an HandleCopyResult()
still having the responsibility to open/close the output stream.
I think this code should be moved upper in the call stack, in
ExecQueryAndProcessResults().

The first attached patch implements a fix that way.

When testing this I've stumbled on another issue nearby: COPY TO
STDOUT followed by \watch should normally produce the error message
"\watch cannot be used with COPY", but the execution goes into a
infinite busy loop instead.
This is because ClearOrSaveAllResults() loops over PQgetResult() until
it returns NULL, but as it turns out, that never happens: it seems
stuck on a PGRES_COPY_OUT result.

While looking to fix that, it occurred to me that it would be
simpler to allow \watch to deal with COPY results rather than
continuing to disallow it. ISTM that before v15, the reason
why PSQLexecWatch() did not want to deal with COPY was to not
bother with a niche use case, rather than because of some
specific impossibility with it.
Now that it calls the generic ExecQueryAndProcessResults() code
that can handle COPY transfers, \watch on a COPY query seems to work
fine if not disallowed.
Besides, v15 adds the possibility to feed \watch output into
a program through PSQL_WATCH_PAGER, and since the copy format is
the best format to be consumed by programs, this seems like
a good reason to allow COPY TO STDOUT with it.
\watch on a COPY FROM STDIN query doesn't make much sense,
but it can be interrupted with ^C if run by mistake, so I don't see a
need to disallow it specifically.

So the second patch fixes the infinite loop problem like that, on top of
the first patch.



Best regards,
-- 
Daniel Vérité
https://postgresql.verite.pro/
Twitter: @DanielVerite
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index e611e3266d..88f4b159f9 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -669,42 +669,13 @@ PrintQueryTuples(const PGresult *result, const 
printQueryOpt *opt, FILE *printQu
 {
        bool            ok = true;
 
-       /* write output to \g argument, if any */
-       if (pset.gfname)
-       {
-               FILE       *fout;
-               bool            is_pipe;
-
-               if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
-                       return false;
-               if (is_pipe)
-                       disable_sigpipe_trap();
-
-               printQuery(result, &pset.popt, fout, false, pset.logfile);
-               if (ferror(fout))
-               {
-                       pg_log_error("could not print result table: %m");
-                       ok = false;
-               }
+       FILE       *fout = printQueryFout ? printQueryFout : pset.queryFout;
 
-               if (is_pipe)
-               {
-                       pclose(fout);
-                       restore_sigpipe_trap();
-               }
-               else
-                       fclose(fout);
-       }
-       else
+       printQuery(result, opt ? opt : &pset.popt, fout, false, pset.logfile);
+       if (ferror(fout))
        {
-               FILE       *fout = printQueryFout ? printQueryFout : 
pset.queryFout;
-
-               printQuery(result, opt ? opt : &pset.popt, fout, false, 
pset.logfile);
-               if (ferror(fout))
-               {
-                       pg_log_error("could not print result table: %m");
-                       ok = false;
-               }
+               pg_log_error("could not print result table: %m");
+               ok = false;
        }
 
        return ok;
@@ -861,10 +832,9 @@ loop_exit:
  * want if the status line doesn't get taken as part of the COPY data.
  */
 static bool
-HandleCopyResult(PGresult **resultp)
+HandleCopyResult(PGresult **resultp, FILE *copystream)
 {
        bool            success;
-       FILE       *copystream;
        PGresult   *copy_result;
        ExecStatusType result_status = PQresultStatus(*resultp);
 
@@ -875,33 +845,6 @@ HandleCopyResult(PGresult **resultp)
 
        if (result_status == PGRES_COPY_OUT)
        {
-               bool            need_close = false;
-               bool            is_pipe = false;
-
-               if (pset.copyStream)
-               {
-                       /* invoked by \copy */
-                       copystream = pset.copyStream;
-               }
-               else if (pset.gfname)
-               {
-                       /* invoked by \g */
-                       if (openQueryOutputFile(pset.gfname,
-                                                                       
&copystream, &is_pipe))
-                       {
-                               need_close = true;
-                               if (is_pipe)
-                                       disable_sigpipe_trap();
-                       }
-                       else
-                               copystream = NULL;      /* discard COPY data 
entirely */
-               }
-               else
-               {
-                       /* fall back to the generic query output stream */
-                       copystream = pset.queryFout;
-               }
-
                success = handleCopyOut(pset.db,
                                                                copystream,
                                                                &copy_result)
@@ -917,24 +860,11 @@ HandleCopyResult(PGresult **resultp)
                        PQclear(copy_result);
                        copy_result = NULL;
                }
-
-               if (need_close)
-               {
-                       /* close \g argument file/pipe */
-                       if (is_pipe)
-                       {
-                               pclose(copystream);
-                               restore_sigpipe_trap();
-                       }
-                       else
-                       {
-                               fclose(copystream);
-                       }
-               }
        }
        else
        {
                /* COPY IN */
+               /* Ignore the copystream argument passed to the function */
                copystream = pset.copyStream ? pset.copyStream : 
pset.cur_cmd_source;
                success = handleCopyIn(pset.db,
                                                           copystream,
@@ -1445,6 +1375,8 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
        instr_time      before,
                                after;
        PGresult   *result;
+       bool            is_pipe = false;
+       FILE            *results_fout = NULL;
 
        if (timing)
                INSTR_TIME_SET_CURRENT(before);
@@ -1555,6 +1487,8 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
                if (result_status == PGRES_COPY_IN ||
                        result_status == PGRES_COPY_OUT)
                {
+                       FILE *copy_stream = NULL;
+
                        if (is_watch)
                        {
                                ClearOrSaveAllResults();
@@ -1562,7 +1496,36 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
                                return -1;
                        }
 
-                       success &= HandleCopyResult(&result);
+                       if (pset.copyStream) /* invoked by \copy */
+                       {
+                               copy_stream = pset.copyStream;
+                       }
+                       else if (pset.gfname)
+                       {
+                               if (results_fout==NULL)
+                               {
+                                       if (openQueryOutputFile(pset.gfname, 
&results_fout, &is_pipe))
+                                       {
+                                               if (is_pipe)
+                                                       disable_sigpipe_trap();
+                                               copy_stream = results_fout;
+                                       }
+                                       else
+                                               success = false;
+                               }
+                               else
+                                       copy_stream = results_fout;
+                       }
+                       else
+                       {
+                               /* fall back to the generic query output stream 
*/
+                               copy_stream = pset.queryFout;
+                       }
+
+                       /* Even if the output stream could not be opened, call
+                          HandleCopyResult() with a NULL output stream to 
collect the
+                          COPY data. */
+                       success &= HandleCopyResult(&result, copy_stream);
                }
 
                /*
@@ -1594,7 +1557,24 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
 
                /* this may or may not print something depending on settings */
                if (result != NULL)
-                       success &= PrintQueryResult(result, last, false, opt, 
printQueryFout);
+               {
+                       /* if results need to be printed into the file 
specified by
+                          \g, open it */
+                       if (PQresultStatus(result) == PGRES_TUPLES_OK &&
+                               pset.gfname && results_fout==NULL)
+                       {
+                               if (openQueryOutputFile(pset.gfname, 
&results_fout, &is_pipe))
+                               {
+                                       if (is_pipe)
+                                               disable_sigpipe_trap();
+                                       printQueryFout = results_fout;
+                               }
+                               else
+                                       success = false;
+                       }
+                       if (success)
+                               success &= PrintQueryResult(result, last, 
false, opt, printQueryFout);
+               }
 
                /* set variables on last result if all went well */
                if (!is_watch && last && success)
@@ -1610,6 +1590,17 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
                }
        }
 
+       if (results_fout)
+       {
+               if (is_pipe)
+               {
+                       pclose(results_fout);
+                       restore_sigpipe_trap();
+               }
+               else
+                       fclose(results_fout);
+       }
+
        /* may need this to recover from conn loss during COPY */
        if (!CheckConnection())
                return -1;
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 88f4b159f9..5864a9833f 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1351,10 +1351,9 @@ DescribeQuery(const char *query, double *elapsed_msec)
  *
  * Sends query and cycles through PGresult objects.
  *
- * When not under \watch and if our command string contained a COPY FROM STDIN
- * or COPY TO STDOUT, the PGresult associated with these commands must be
- * processed by providing an input or output stream.  In that event, we'll
- * marshal data for the COPY.
+ * If our command string contained a COPY FROM STDIN or COPY TO STDOUT, the
+ * PGresult associated with these commands must be processed by providing an
+ * input or output stream.  In that event, we'll marshal data for the COPY.
  *
  * For other commands, the results are processed normally, depending on their
  * status.
@@ -1491,12 +1490,9 @@ ExecQueryAndProcessResults(const char *query, double 
*elapsed_msec, bool *svpt_g
 
                        if (is_watch)
                        {
-                               ClearOrSaveAllResults();
-                               pg_log_error("\\watch cannot be used with 
COPY");
-                               return -1;
+                               copy_stream = printQueryFout ? printQueryFout : 
pset.queryFout;
                        }
-
-                       if (pset.copyStream) /* invoked by \copy */
+                       else if (pset.copyStream) /* invoked by \copy */
                        {
                                copy_stream = pset.copyStream;
                        }

Reply via email to