diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/bin/psql/common.c | 52 | ||||
| -rw-r--r-- | src/bin/psql/copy.c | 22 | ||||
| -rw-r--r-- | src/bin/psql/copy.h | 8 |
3 files changed, 53 insertions, 29 deletions
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 6ca9bbc9d86..6968adfd42c 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -632,7 +632,9 @@ StoreQueryTuple(const PGresult *result) * degenerates to an AcceptResult() call. * * Changes its argument to point to the last PGresult of the command string, - * or NULL if that result was for a COPY FROM STDIN or COPY TO STDOUT. + * or NULL if that result was for a COPY TO STDOUT. (Returning NULL prevents + * the command status from being printed, which we want in that case so that + * the status line doesn't get taken as part of the COPY data.) * * Returns true on complete success, false otherwise. Possible failure modes * include purely client-side problems; check the transaction status for the @@ -641,14 +643,14 @@ StoreQueryTuple(const PGresult *result) static bool ProcessResult(PGresult **results) { - PGresult *next_result; bool success = true; bool first_cycle = true; - do + for (;;) { ExecStatusType result_status; bool is_copy; + PGresult *next_result; if (!AcceptResult(*results)) { @@ -693,6 +695,7 @@ ProcessResult(PGresult **results) * otherwise use queryFout or cur_cmd_source as appropriate. */ FILE *copystream = pset.copyStream; + PGresult *copy_result; SetCancelConn(); if (result_status == PGRES_COPY_OUT) @@ -700,7 +703,19 @@ ProcessResult(PGresult **results) if (!copystream) copystream = pset.queryFout; success = handleCopyOut(pset.db, - copystream) && success; + copystream, + ©_result) && success; + + /* + * 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; + } } else { @@ -708,30 +723,37 @@ ProcessResult(PGresult **results) copystream = pset.cur_cmd_source; success = handleCopyIn(pset.db, copystream, - PQbinaryTuples(*results)) && success; + PQbinaryTuples(*results), + ©_result) && success; } ResetCancelConn(); /* - * Call PQgetResult() once more. In the typical case of a - * single-command string, it will return NULL. Otherwise, we'll - * have other results to process that may include other COPYs. + * Replace the PGRES_COPY_OUT/IN result with COPY command's exit + * status, or with NULL if we want to suppress printing anything. */ PQclear(*results); - *results = next_result = PQgetResult(pset.db); + *results = copy_result; } else if (first_cycle) + { /* fast path: no COPY commands; PQexec visited all results */ break; - else if ((next_result = PQgetResult(pset.db))) - { - /* non-COPY command(s) after a COPY: keep the last one */ - PQclear(*results); - *results = next_result; } + /* + * Check PQgetResult() again. In the typical case of a single-command + * string, it will return NULL. Otherwise, we'll have other results + * to process that may include other COPYs. We keep the last result. + */ + next_result = PQgetResult(pset.db); + if (!next_result) + break; + + PQclear(*results); + *results = next_result; first_cycle = false; - } while (next_result); + } /* may need this to recover from conn loss during COPY */ if (!first_cycle && !CheckConnection()) diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c index a058f2ff0d4..d706206cc96 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -429,16 +429,17 @@ do_copy(const char *args) * conn should be a database connection that you just issued COPY TO on * and got back a PGRES_COPY_OUT result. * copystream is the file stream for the data to go to. + * The final status for the COPY is returned into *res (but note + * we already reported the error, if it's not a success result). * * result is true if successful, false if not. */ bool -handleCopyOut(PGconn *conn, FILE *copystream) +handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res) { bool OK = true; char *buf; int ret; - PGresult *res; for (;;) { @@ -485,13 +486,12 @@ handleCopyOut(PGconn *conn, FILE *copystream) * but hasn't exited COPY_OUT state internally. So we ignore the * possibility here. */ - res = PQgetResult(conn); - if (PQresultStatus(res) != PGRES_COMMAND_OK) + *res = PQgetResult(conn); + if (PQresultStatus(*res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); OK = false; } - PQclear(res); return OK; } @@ -504,6 +504,8 @@ handleCopyOut(PGconn *conn, FILE *copystream) * and got back a PGRES_COPY_IN result. * copystream is the file stream to read the data from. * isbinary can be set from PQbinaryTuples(). + * The final status for the COPY is returned into *res (but note + * we already reported the error, if it's not a success result). * * result is true if successful, false if not. */ @@ -512,12 +514,11 @@ handleCopyOut(PGconn *conn, FILE *copystream) #define COPYBUFSIZ 8192 bool -handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary) +handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) { bool OK; const char *prompt; char buf[COPYBUFSIZ]; - PGresult *res; /* * Establish longjmp destination for exiting from wait-for-input. (This is @@ -679,21 +680,20 @@ copyin_cleanup: * connection is lost. But that's fine; it will get us out of COPY_IN * state, which is what we need.) */ - while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_IN) + while (*res = PQgetResult(conn), PQresultStatus(*res) == PGRES_COPY_IN) { OK = false; - PQclear(res); + PQclear(*res); /* We can't send an error message if we're using protocol version 2 */ PQputCopyEnd(conn, (PQprotocolVersion(conn) < 3) ? NULL : _("trying to exit copy mode")); } - if (PQresultStatus(res) != PGRES_COMMAND_OK) + if (PQresultStatus(*res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); OK = false; } - PQclear(res); return OK; } diff --git a/src/bin/psql/copy.h b/src/bin/psql/copy.h index ec1f0d09a9a..2c71da00603 100644 --- a/src/bin/psql/copy.h +++ b/src/bin/psql/copy.h @@ -12,11 +12,13 @@ /* handler for \copy */ -bool do_copy(const char *args); +extern bool do_copy(const char *args); /* lower level processors for copy in/out streams */ -bool handleCopyOut(PGconn *conn, FILE *copystream); -bool handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary); +extern bool handleCopyOut(PGconn *conn, FILE *copystream, + PGresult **res); +extern bool handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, + PGresult **res); #endif |