psql: Refactor ProcessResult()

Separate HandleCopyResult() from ProcessResult() in preparation for a
subsequent patch.

Author: Fabien COELHO <coelho@cri.ensmp.fr>
Discussion: https://www.postgresql.org/message-id/flat/alpine.DEB.2.21.1904132231510.8961@lancre
This commit is contained in:
Peter Eisentraut 2022-04-01 13:00:43 +02:00
parent d16773cdc8
commit 5519d5affd
1 changed files with 111 additions and 93 deletions

View File

@ -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,
&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)
&& (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),
&copy_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,
&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)
&& 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),
&copy_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 */