diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index df8c80a1b26..7fc466a3026 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -639,7 +639,7 @@ PrintQueryTuples(const PGresult *results) * 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, @@ -699,13 +699,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 7459a8c7c87..8258f0592ca 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -244,8 +244,6 @@ do_copy(const char *args) { PQExpBufferData query; FILE *copystream; - FILE *save_file; - FILE **override_file; struct copy_options *options; bool success; struct stat st; @@ -262,8 +260,6 @@ do_copy(const char *args) if (options->from) { - override_file = &pset.cur_cmd_source; - if (options->file) copystream = fopen(options->file, PG_BINARY_R); else if (!options->psql_inout) @@ -273,8 +269,6 @@ do_copy(const char *args) } else { - override_file = &pset.queryFout; - if (options->file) copystream = fopen(options->file, PG_BINARY_W); else if (!options->psql_inout) @@ -313,11 +307,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 c907fa0003a..bfd951a2804 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 b30639395ec..94e9ab17ebf 100644 --- a/src/bin/psql/startup.c +++ b/src/bin/psql/startup.c @@ -120,6 +120,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;