diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index c9847c8f9a..f6777bbfc4 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -892,6 +892,116 @@ loop_exit: } +/* + * Marshal the COPY data. Either subroutine 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 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. + * 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) +{ + bool success; + FILE *copystream; + PGresult *copy_result; + ExecStatusType result_status = PQresultStatus(*resultp); + + Assert(result_status == PGRES_COPY_OUT || + result_status == PGRES_COPY_IN); + + SetCancelConn(pset.db); + + 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) + && (copystream != NULL); + + /* + * Suppress status printing if the report would go to the same + * place as the COPY data just went. Note this doesn't + * prevent error reporting, since handleCopyOut did that. + */ + if (copystream == pset.queryFout) + { + 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 */ + copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source; + success = handleCopyIn(pset.db, + copystream, + PQbinaryTuples(*resultp), + ©_result); + } + ResetCancelConn(); + + /* + * Replace the PGRES_COPY_OUT/IN result with COPY command's exit + * status, or with NULL if we want to suppress printing anything. + */ + PQclear(*resultp); + *resultp = copy_result; + + return success; +} + + /* * ProcessResult: utility function for use by SendQuery() only * @@ -957,99 +1067,7 @@ ProcessResult(PGresult **resultp) } if (is_copy) - { - /* - * Marshal the COPY data. Either subroutine will get the - * connection out of its COPY state, then call PQresultStatus() - * once and report any error. - * - * 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 IN, use pset.copyStream as data source if it's set, - * otherwise cur_cmd_source. - */ - FILE *copystream; - PGresult *copy_result; - - SetCancelConn(pset.db); - 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) - && success - && (copystream != NULL); - - /* - * Suppress status printing if the report would go to the same - * place as the COPY data just went. Note this doesn't - * prevent error reporting, since handleCopyOut did that. - */ - if (copystream == pset.queryFout) - { - 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 */ - copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source; - success = handleCopyIn(pset.db, - copystream, - PQbinaryTuples(*resultp), - ©_result) && success; - } - ResetCancelConn(); - - /* - * Replace the PGRES_COPY_OUT/IN result with COPY command's exit - * status, or with NULL if we want to suppress printing anything. - */ - PQclear(*resultp); - *resultp = copy_result; - } + success = HandleCopyResult(resultp); else if (first_cycle) { /* fast path: no COPY commands; PQexec visited all results */