diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 3a820fa..136eed1 100644 *** a/src/bin/psql/common.c --- b/src/bin/psql/common.c *************** StoreQueryTuple(const PGresult *result) *** 628,638 **** * command. In that event, we'll marshal data for the COPY and then cycle * through any subsequent PGresult objects. * ! * When the command string contained no affected COPY command, this function * 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. * * Returns true on complete success, false otherwise. Possible failure modes * include purely client-side problems; check the transaction status for the --- 628,637 ---- * command. In that event, we'll marshal data for the COPY and then cycle * through any subsequent PGresult objects. * ! * When the command string contained no such COPY command, this function * degenerates to an AcceptResult() call. * ! * Changes its argument to point to the last PGresult of the command string. * * Returns true on complete success, false otherwise. Possible failure modes * include purely client-side problems; check the transaction status for the *************** StoreQueryTuple(const PGresult *result) *** 641,654 **** static bool ProcessResult(PGresult **results) { - PGresult *next_result; bool success = true; bool first_cycle = true; ! do { ExecStatusType result_status; bool is_copy; if (!AcceptResult(*results)) { --- 640,653 ---- static bool ProcessResult(PGresult **results) { bool success = true; bool first_cycle = true; ! for (;;) { ExecStatusType result_status; bool is_copy; + PGresult *next_result; if (!AcceptResult(*results)) { *************** ProcessResult(PGresult **results) *** 688,722 **** * Marshal the COPY data. Either subroutine will get the * connection out of its COPY state, then call PQresultStatus() * once and report any error. */ SetCancelConn(); if (result_status == PGRES_COPY_OUT) ! success = handleCopyOut(pset.db, pset.queryFout) && success; else ! success = handleCopyIn(pset.db, pset.cur_cmd_source, ! PQbinaryTuples(*results)) && 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. ! */ PQclear(*results); ! *results = next_result = PQgetResult(pset.db); } 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; } first_cycle = false; ! } while (next_result); /* may need this to recover from conn loss during COPY */ if (!first_cycle && !CheckConnection()) --- 687,742 ---- * Marshal the COPY data. Either subroutine will get the * connection out of its COPY state, then call PQresultStatus() * once and report any error. + * + * If pset.copyStream is set, use that as data source/sink, + * otherwise use queryFout or cur_cmd_source as appropriate. */ + FILE *copystream = pset.copyStream; + PGresult *copy_result; + SetCancelConn(); if (result_status == PGRES_COPY_OUT) ! { ! if (!copystream) ! copystream = pset.queryFout; ! success = handleCopyOut(pset.db, ! copystream, ! ©_result) && success; ! } else ! { ! if (!copystream) ! copystream = pset.cur_cmd_source; ! success = handleCopyIn(pset.db, ! copystream, ! PQbinaryTuples(*results), ! ©_result) && success; ! } ResetCancelConn(); ! /* replace the COPY_OUT/IN result with COPY command exit status */ PQclear(*results); ! *results = copy_result; } else if (first_cycle) + { /* fast path: no COPY commands; PQexec visited all results */ break; } + /* + * 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; ! } /* 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 9e815b1..d706206 100644 *** a/src/bin/psql/copy.c --- b/src/bin/psql/copy.c *************** do_copy(const char *args) *** 269,279 **** { PQExpBufferData query; FILE *copystream; - FILE *save_file; - FILE **override_file; struct copy_options *options; bool success; - struct stat st; /* parse options */ options = parse_slash_copy(args); --- 269,276 ---- *************** do_copy(const char *args) *** 287,294 **** if (options->from) { - override_file = &pset.cur_cmd_source; - if (options->file) { if (options->program) --- 284,289 ---- *************** do_copy(const char *args) *** 308,315 **** } else { - override_file = &pset.queryFout; - if (options->file) { if (options->program) --- 303,308 ---- *************** do_copy(const char *args) *** 345,350 **** --- 338,344 ---- if (!options->program) { + struct stat st; int result; /* make sure the specified file is not a directory */ *************** do_copy(const char *args) *** 375,385 **** if (options->after_tofrom) appendPQExpBufferStr(&query, options->after_tofrom); ! /* Run it like a user command, interposing the data source or sink. */ ! save_file = *override_file; ! *override_file = copystream; success = SendQuery(query.data); ! *override_file = save_file; termPQExpBuffer(&query); if (options->file != NULL) --- 369,378 ---- if (options->after_tofrom) appendPQExpBufferStr(&query, options->after_tofrom); ! /* run it like a user command, but with copystream as data source/sink */ ! pset.copyStream = copystream; success = SendQuery(query.data); ! pset.copyStream = NULL; termPQExpBuffer(&query); if (options->file != NULL) *************** do_copy(const char *args) *** 436,451 **** * 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. * * result is true if successful, false if not. */ bool ! handleCopyOut(PGconn *conn, FILE *copystream) { bool OK = true; char *buf; int ret; - PGresult *res; for (;;) { --- 429,445 ---- * 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, PGresult **res) { bool OK = true; char *buf; int ret; for (;;) { *************** handleCopyOut(PGconn *conn, FILE *copyst *** 492,504 **** * but hasn't exited COPY_OUT state internally. So we ignore the * possibility here. */ ! res = PQgetResult(conn); ! if (PQresultStatus(res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); OK = false; } - PQclear(res); return OK; } --- 486,497 ---- * but hasn't exited COPY_OUT state internally. So we ignore the * possibility here. */ ! *res = PQgetResult(conn); ! if (PQresultStatus(*res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); OK = false; } return OK; } *************** handleCopyOut(PGconn *conn, FILE *copyst *** 511,516 **** --- 504,511 ---- * 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. */ *************** handleCopyOut(PGconn *conn, FILE *copyst *** 519,530 **** #define COPYBUFSIZ 8192 bool ! handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary) { bool OK; const char *prompt; char buf[COPYBUFSIZ]; - PGresult *res; /* * Establish longjmp destination for exiting from wait-for-input. (This is --- 514,524 ---- #define COPYBUFSIZ 8192 bool ! handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) { bool OK; const char *prompt; char buf[COPYBUFSIZ]; /* * Establish longjmp destination for exiting from wait-for-input. (This is *************** copyin_cleanup: *** 686,706 **** * 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) { OK = false; ! 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) { psql_error("%s", PQerrorMessage(conn)); OK = false; } - PQclear(res); return OK; } --- 680,699 ---- * 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) { OK = false; ! 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) { psql_error("%s", PQerrorMessage(conn)); OK = false; } return OK; } diff --git a/src/bin/psql/copy.h b/src/bin/psql/copy.h index ec1f0d0..2c71da0 100644 *** a/src/bin/psql/copy.h --- b/src/bin/psql/copy.h *************** *** 12,22 **** /* handler for \copy */ ! 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); #endif --- 12,24 ---- /* handler for \copy */ ! extern bool do_copy(const char *args); /* lower level processors for copy in/out streams */ ! extern bool handleCopyOut(PGconn *conn, FILE *copystream, ! PGresult **res); ! extern bool handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, ! PGresult **res); #endif diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h index 3e8328d..eecffb1 100644 *** a/src/bin/psql/settings.h --- b/src/bin/psql/settings.h *************** typedef struct _psqlSettings *** 70,75 **** --- 70,77 ---- FILE *queryFout; /* where to send the query results */ bool queryFoutPipe; /* queryFout is from a popen() */ + FILE *copyStream; /* Stream to read/write for \copy command */ + printQueryOpt popt; char *gfname; /* one-shot file output argument for \g */ diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c index d5f1c0d..45653a1 100644 *** a/src/bin/psql/startup.c --- b/src/bin/psql/startup.c *************** main(int argc, char *argv[]) *** 118,123 **** --- 118,124 ---- pset.encoding = PQenv2encoding(); pset.queryFout = stdout; pset.queryFoutPipe = false; + pset.copyStream = NULL; pset.cur_cmd_source = stdin; pset.cur_cmd_interactive = false;