diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index 0dea54d..7398dec 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -65,6 +65,8 @@ static backslashResult exec_command_cd(PsqlScanState scan_state, bool active_bra const char *cmd); static backslashResult exec_command_conninfo(PsqlScanState scan_state, bool active_branch); static backslashResult exec_command_copy(PsqlScanState scan_state, bool active_branch); +static backslashResult exec_command_copyfrom(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch); +static backslashResult exec_command_copyto(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch); static backslashResult exec_command_copyright(PsqlScanState scan_state, bool active_branch); static backslashResult exec_command_crosstabview(PsqlScanState scan_state, bool active_branch); static backslashResult exec_command_d(PsqlScanState scan_state, bool active_branch, @@ -305,6 +307,10 @@ exec_command(const char *cmd, status = exec_command_conninfo(scan_state, active_branch); else if (pg_strcasecmp(cmd, "copy") == 0) status = exec_command_copy(scan_state, active_branch); + else if (pg_strcasecmp(cmd, "copyfrom") == 0) + status = exec_command_copyfrom(query_buf, scan_state, active_branch); + else if (pg_strcasecmp(cmd, "copyto") == 0) + status = exec_command_copyto(query_buf, scan_state, active_branch); else if (strcmp(cmd, "copyright") == 0) status = exec_command_copyright(scan_state, active_branch); else if (strcmp(cmd, "crosstabview") == 0) @@ -634,6 +640,53 @@ exec_command_copy(PsqlScanState scan_state, bool active_branch) } /* + * \copyfrom -- run a COPY FROM command + * The COPY query is obtained from the query buffer + * The argument is 'filename' as in \copy + */ +static backslashResult +exec_command_copyfrom(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch) +{ + bool success = true; + + if (active_branch) + { + char *opt = psql_scan_slash_option(scan_state, + OT_FILEPIPE, NULL, false); + + success = do_copy_query(query_buf->data, opt, copy_from); + resetPQExpBuffer(query_buf); + free(opt); + } + else + ignore_slash_filepipe(scan_state); + + return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR; +} +/* + * \copyto -- run a COPY TO command + */ +static backslashResult +exec_command_copyto(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch) +{ + bool success = true; + + if (active_branch) + { + char *opt = psql_scan_slash_option(scan_state, + OT_FILEPIPE, NULL, false); + + success = do_copy_query(query_buf->data, opt, copy_to); + resetPQExpBuffer(query_buf); + free(opt); + } + else + ignore_slash_filepipe(scan_state); + + return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR; +} + +/* * \copyright -- print copyright notice */ static backslashResult diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c index 555c633..95c5c5e 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -260,28 +260,11 @@ error: } -/* - * Execute a \copy command (frontend copy). We have to open a file (or execute - * a command), then submit a COPY query to the backend and either feed it data - * from the file or route its response into the file. - */ -bool -do_copy(const char *args) +static +FILE* +open_copy_stream(struct copy_options *options) { - PQExpBufferData query; - FILE *copystream; - struct copy_options *options; - bool success; - - /* parse options */ - options = parse_slash_copy(args); - - if (!options) - return false; - - /* prepare to read or write the target file */ - if (options->file && !options->program) - canonicalize_path(options->file); + FILE *copystream = NULL; if (options->from) { @@ -331,8 +314,7 @@ do_copy(const char *args) else psql_error("%s: %s\n", options->file, strerror(errno)); - free_copy_options(options); - return false; + return NULL; } if (!options->program) @@ -352,27 +334,18 @@ do_copy(const char *args) if (result < 0 || S_ISDIR(st.st_mode)) { fclose(copystream); - free_copy_options(options); - return false; + return NULL; } } - /* build the command we will send to the backend */ - initPQExpBuffer(&query); - printfPQExpBuffer(&query, "COPY "); - appendPQExpBufferStr(&query, options->before_tofrom); - if (options->from) - appendPQExpBufferStr(&query, " FROM STDIN "); - else - appendPQExpBufferStr(&query, " TO STDOUT "); - if (options->after_tofrom) - appendPQExpBufferStr(&query, options->after_tofrom); + return copystream; +} - /* 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); +static +bool +close_copy_stream(FILE* copystream, struct copy_options *options) +{ + bool success = true; if (options->file != NULL) { @@ -407,6 +380,119 @@ do_copy(const char *args) } } } + return success; +} + +/* + * Execute a \copy command (frontend copy). We have to open a file (or execute + * a command), then submit a COPY query to the backend and either feed it data + * from the file or route its response into the file. + */ +bool +do_copy(const char *args) +{ + PQExpBufferData query; + FILE *copystream; + struct copy_options *options; + bool success; + + /* parse options */ + options = parse_slash_copy(args); + + if (!options) + return false; + + /* prepare to read or write the target file */ + if (options->file && !options->program) + canonicalize_path(options->file); + + /* open the file or program from which or to which data goes */ + copystream = open_copy_stream(options); + + if (!copystream) + { + /* a user-visible error has already been emitted at this point, so + * just clean up and return */ + free_copy_options(options); + return false; + } + + /* build the command we will send to the backend */ + initPQExpBuffer(&query); + printfPQExpBuffer(&query, "COPY "); + appendPQExpBufferStr(&query, options->before_tofrom); + if (options->from) + appendPQExpBufferStr(&query, " FROM STDIN "); + else + appendPQExpBufferStr(&query, " TO STDOUT "); + 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); + success = close_copy_stream(copystream, options); + free_copy_options(options); + return success; +} + +/* + * Send the COPY query (frontend copy) in the query buffer, and feed it data + * from the file or program given as argument (when dir=copy_from), + * or route its response into the file or program (when dir=copy_to). + */ +bool +do_copy_query(const char* querystring, const char *args, enum copy_direction dir) +{ + FILE *copystream; + struct copy_options *options; + bool success; + + options = pg_malloc0(sizeof(struct copy_options)); + if (!options) + return false; + + options->from = (dir == copy_from); + + if (args) + { + if (args[0] == '|') + { + options->program = true; + options->file = pg_strdup(args+1); + } + else + { + options->file = pg_strdup(args); + canonicalize_path(options->file); + } + } + else + { + /* Use psql's stdin or stdout if no file or program given */ + options->psql_inout = true; + } + + /* open the file or program from which or to which data goes */ + copystream = open_copy_stream(options); + if (!copystream) + { + /* a user-visible error has already been emitted at this point, so + * just clean up and return */ + free_copy_options(options); + return false; + } + + /* run the COPY query with copystream as data source/sink */ + pset.copyStream = copystream; + success = SendQuery(querystring); + pset.copyStream = NULL; + + success = close_copy_stream(copystream, options); + free_copy_options(options); return success; } diff --git a/src/bin/psql/copy.h b/src/bin/psql/copy.h index f4107d7..7755ac7 100644 --- a/src/bin/psql/copy.h +++ b/src/bin/psql/copy.h @@ -10,10 +10,18 @@ #include "libpq-fe.h" +enum copy_direction +{ + copy_from = 1, + copy_to +}; /* handler for \copy */ extern bool do_copy(const char *args); +/* handler for \copyfrom and \copyto */ +extern bool do_copy_query(const char* querystring, const char *args, enum copy_direction dir); + /* lower level processors for copy in/out streams */ extern bool handleCopyOut(PGconn *conn, FILE *copystream,