diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 974959c5959..53e710e92f9 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -32,8 +32,12 @@ static bool DescribeQuery(const char *query, double *elapsed_msec); static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec); -static int ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_gone_p, - bool is_watch, const printQueryOpt *opt, FILE *printQueryFout); +static int ExecQueryAndProcessResults(const char *query, + double *elapsed_msec, + bool *svpt_gone_p, + bool is_watch, + const printQueryOpt *opt, + FILE *printQueryFout); static bool command_no_begin(const char *query); static bool is_select_command(const char *query); @@ -662,49 +666,27 @@ PrintNotifications(void) /* * PrintQueryTuples: assuming query result is OK, print its tuples * + * We use the options given by opt unless that's NULL, in which case + * we use pset.popt. + * + * Output is to printQueryFout unless that's NULL, in which case + * we use pset.queryFout. + * * Returns true if successful, false otherwise. */ static bool -PrintQueryTuples(const PGresult *result, const printQueryOpt *opt, FILE *printQueryFout) +PrintQueryTuples(const PGresult *result, const printQueryOpt *opt, + FILE *printQueryFout) { bool ok = true; + FILE *fout = printQueryFout ? printQueryFout : pset.queryFout; - /* write output to \g argument, if any */ - if (pset.gfname) + printQuery(result, opt ? opt : &pset.popt, fout, false, pset.logfile); + fflush(fout); + if (ferror(fout)) { - 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; - } - - if (is_pipe) - { - pclose(fout); - restore_sigpipe_trap(); - } - else - fclose(fout); - } - else - { - 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; @@ -845,26 +827,24 @@ loop_exit: /* - * Marshal the COPY data. Either subroutine will get the + * Marshal the COPY data. Either path will get the * connection out of its COPY state, then call PQresultStatus() * once and report any error. Return whether all was ok. * - * For COPY OUT, direct the output to pset.copyStream if it's set, - * otherwise to pset.gfname if it's set, otherwise to queryFout. + * For COPY OUT, direct the output to copystream, or discard if that's NULL. * For COPY IN, use pset.copyStream as data source if it's set, * otherwise cur_cmd_source. * - * Update result if further processing is necessary, or NULL otherwise. + * Update *resultp if further processing is necessary; set to NULL otherwise. * Return a result when queryFout can safely output a result status: on COPY * IN, or on COPY OUT if written to something other than pset.queryFout. * Returning NULL prevents the command status from being printed, which we * 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 +855,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, - ©stream, &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, ©_result) @@ -917,24 +870,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, @@ -974,6 +914,7 @@ PrintQueryStatus(PGresult *result, FILE *printQueryFout) } else fprintf(fout, "%s\n", PQcmdStatus(result)); + fflush(fout); } if (pset.logfile) @@ -989,10 +930,16 @@ PrintQueryStatus(PGresult *result, FILE *printQueryFout) * * Note: Utility function for use by SendQuery() only. * + * last is true if this is the last result of a command string. + * opt and printQueryFout are defined as for PrintQueryTuples. + * printStatusFout is where to send command status; NULL means pset.queryFout. + * * Returns true if the query executed successfully, false otherwise. */ static bool -PrintQueryResult(PGresult *result, bool last, bool is_watch, const printQueryOpt *opt, FILE *printQueryFout) +PrintQueryResult(PGresult *result, bool last, + const printQueryOpt *opt, FILE *printQueryFout, + FILE *printStatusFout) { bool success; const char *cmdstatus; @@ -1022,14 +969,14 @@ PrintQueryResult(PGresult *result, bool last, bool is_watch, const printQueryOpt if (strncmp(cmdstatus, "INSERT", 6) == 0 || strncmp(cmdstatus, "UPDATE", 6) == 0 || strncmp(cmdstatus, "DELETE", 6) == 0) - PrintQueryStatus(result, printQueryFout); + PrintQueryStatus(result, printStatusFout); } break; case PGRES_COMMAND_OK: if (last || pset.show_all_results) - PrintQueryStatus(result, printQueryFout); + PrintQueryStatus(result, printStatusFout); success = true; break; @@ -1056,8 +1003,6 @@ PrintQueryResult(PGresult *result, bool last, bool is_watch, const printQueryOpt break; } - fflush(printQueryFout ? printQueryFout : pset.queryFout); - return success; } @@ -1399,7 +1344,7 @@ DescribeQuery(const char *query, double *elapsed_msec) } if (OK && result) - OK = PrintQueryResult(result, true, false, NULL, NULL); + OK = PrintQueryResult(result, true, NULL, NULL, NULL); termPQExpBuffer(&buf); } @@ -1421,10 +1366,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. @@ -1437,14 +1381,18 @@ DescribeQuery(const char *query, double *elapsed_msec) * committed. */ static int -ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_gone_p, - bool is_watch, const printQueryOpt *opt, FILE *printQueryFout) +ExecQueryAndProcessResults(const char *query, + double *elapsed_msec, bool *svpt_gone_p, + bool is_watch, + const printQueryOpt *opt, FILE *printQueryFout) { bool timing = pset.timing; bool success; instr_time before, after; PGresult *result; + FILE *gfile_fout = NULL; + bool gfile_is_pipe = false; if (timing) INSTR_TIME_SET_CURRENT(before); @@ -1555,14 +1503,57 @@ ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_g if (result_status == PGRES_COPY_IN || result_status == PGRES_COPY_OUT) { - if (is_watch) + FILE *copy_stream = NULL; + + /* + * For COPY OUT, direct the output to the default place (probably + * a pager pipe) for \watch, or to pset.copyStream for \copy, + * otherwise to pset.gfname if that's set, otherwise to + * pset.queryFout. + */ + if (result_status == PGRES_COPY_OUT) { - ClearOrSaveAllResults(); - pg_log_error("\\watch cannot be used with COPY"); - return -1; + if (is_watch) + { + /* invoked by \watch */ + copy_stream = printQueryFout ? printQueryFout : pset.queryFout; + } + else if (pset.copyStream) + { + /* invoked by \copy */ + copy_stream = pset.copyStream; + } + else if (pset.gfname) + { + /* send to \g file, which we may have opened already */ + if (gfile_fout == NULL) + { + if (openQueryOutputFile(pset.gfname, + &gfile_fout, &gfile_is_pipe)) + { + if (gfile_is_pipe) + disable_sigpipe_trap(); + copy_stream = gfile_fout; + } + else + success = false; + } + else + copy_stream = gfile_fout; + } + else + { + /* fall back to the generic query output stream */ + copy_stream = pset.queryFout; + } } - success &= HandleCopyResult(&result); + /* + * Even if the output stream could not be opened, we call + * HandleCopyResult() with a NULL output stream to collect and + * discard the COPY data. + */ + success &= HandleCopyResult(&result, copy_stream); } /* @@ -1594,7 +1585,36 @@ 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, unless we already did. Note that when pset.gfname is + * set, the passed-in value of printQueryFout is not used for + * tuple output, but it's still used for status output. + */ + FILE *tuples_fout = printQueryFout; + bool do_print = true; + + if (PQresultStatus(result) == PGRES_TUPLES_OK && + pset.gfname) + { + if (gfile_fout == NULL) + { + if (openQueryOutputFile(pset.gfname, + &gfile_fout, &gfile_is_pipe)) + { + if (gfile_is_pipe) + disable_sigpipe_trap(); + } + else + success = do_print = false; + } + tuples_fout = gfile_fout; + } + if (do_print) + success &= PrintQueryResult(result, last, opt, + tuples_fout, printQueryFout); + } /* set variables on last result if all went well */ if (!is_watch && last && success) @@ -1610,6 +1630,18 @@ ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_g } } + /* close \g file if we opened it */ + if (gfile_fout) + { + if (gfile_is_pipe) + { + pclose(gfile_fout); + restore_sigpipe_trap(); + } + else + fclose(gfile_fout); + } + /* may need this to recover from conn loss during COPY */ if (!CheckConnection()) return -1;