diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 3a820faf78..6ca9bbc9d8 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -628,7 +628,7 @@ StoreQueryTuple(const PGresult *result) * 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 + * 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, @@ -688,13 +688,28 @@ ProcessResult(PGresult **results) * 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; + SetCancelConn(); if (result_status == PGRES_COPY_OUT) - success = handleCopyOut(pset.db, pset.queryFout) && success; + { + if (!copystream) + copystream = pset.queryFout; + success = handleCopyOut(pset.db, + copystream) && success; + } else - success = handleCopyIn(pset.db, pset.cur_cmd_source, + { + if (!copystream) + copystream = pset.cur_cmd_source; + success = handleCopyIn(pset.db, + copystream, PQbinaryTuples(*results)) && success; + } ResetCancelConn(); /* diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c index 9e815b136e..a058f2ff0d 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -269,11 +269,8 @@ do_copy(const char *args) { 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); @@ -287,8 +284,6 @@ do_copy(const char *args) if (options->from) { - override_file = &pset.cur_cmd_source; - if (options->file) { if (options->program) @@ -308,8 +303,6 @@ do_copy(const char *args) } else { - override_file = &pset.queryFout; - if (options->file) { if (options->program) @@ -345,6 +338,7 @@ do_copy(const char *args) if (!options->program) { + struct stat st; int result; /* make sure the specified file is not a directory */ @@ -375,11 +369,10 @@ do_copy(const char *args) 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; + /* run it like a user command, but with copystream as data source/sink */ + pset.copyStream = copystream; success = SendQuery(query.data); - *override_file = save_file; + pset.copyStream = NULL; termPQExpBuffer(&query); if (options->file != NULL) diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h index 3e8328db9a..eecffb1073 100644 --- a/src/bin/psql/settings.h +++ b/src/bin/psql/settings.h @@ -70,6 +70,8 @@ typedef struct _psqlSettings 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 d5f1c0d408..45653a15a8 100644 --- a/src/bin/psql/startup.c +++ b/src/bin/psql/startup.c @@ -118,6 +118,7 @@ main(int argc, char *argv[]) pset.encoding = PQenv2encoding(); pset.queryFout = stdout; pset.queryFoutPipe = false; + pset.copyStream = NULL; pset.cur_cmd_source = stdin; pset.cur_cmd_interactive = false;