diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index 6c3f0aa69d..ab809ec3a0 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -320,25 +320,10 @@ exec_command(const char *cmd, /* \copy */ else if (pg_strcasecmp(cmd, "copy") == 0) { - /* Default fetch-it-all-and-print mode */ - instr_time before, - after; - char *opt = psql_scan_slash_option(scan_state, OT_WHOLE_LINE, NULL, false); - if (pset.timing) - INSTR_TIME_SET_CURRENT(before); - success = do_copy(opt); - - if (pset.timing && success) - { - INSTR_TIME_SET_CURRENT(after); - INSTR_TIME_SUBTRACT(after, before); - printf(_("Time: %.3f ms\n"), INSTR_TIME_GET_MILLISEC(after)); - } - free(opt); } diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 29389d00ad..5c9bd96002 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -438,7 +438,7 @@ ResetCancelConn(void) static bool AcceptResult(const PGresult *result) { - bool OK = true; + bool OK; if (!result) OK = false; @@ -450,11 +450,21 @@ AcceptResult(const PGresult *result) case PGRES_EMPTY_QUERY: case PGRES_COPY_IN: case PGRES_COPY_OUT: + case PGRES_COPY_BOTH: /* Fine, do nothing */ + OK = true; + break; + + case PGRES_BAD_RESPONSE: + case PGRES_NONFATAL_ERROR: + case PGRES_FATAL_ERROR: + OK = false; break; default: OK = false; + psql_error("unexpected PQresultStatus (%d)", + PQresultStatus(result)); break; } @@ -620,45 +630,114 @@ PrintQueryTuples(const PGresult *results) /* - * ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it + * ProcessResult: utility function for use by SendQuery() only * - * Note: Utility function for use by SendQuery() only. + * When our command string contained a COPY FROM STDIN or COPY TO STDOUT, + * PQexec() has stopped at the PGresult associated with the first such + * command. In that event, we'll marshal data for the COPY and then cycle + * through any subsequent PGresult objects. * - * Returns true if the query executed successfully, false otherwise. + * 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 + * server-side opinion. */ static bool -ProcessCopyResult(PGresult *results) +ProcessResult(PGresult **results) { - bool success = false; + PGresult *next_result; + bool success = true; + bool first_cycle = true; - if (!results) - return false; - - switch (PQresultStatus(results)) + do { - case PGRES_TUPLES_OK: - case PGRES_COMMAND_OK: - case PGRES_EMPTY_QUERY: - /* nothing to do here */ - success = true; - break; + ExecStatusType result_status; + bool is_copy; - case PGRES_COPY_OUT: + if (!AcceptResult(*results)) + { + /* + * Failure at this point is always a server-side failure or a + * failure to submit the command string. Either way, we're + * finished with this command string. + */ + success = false; + break; + } + + result_status = PQresultStatus(*results); + switch (result_status) + { + case PGRES_COPY_BOTH: + /* + * No now-existing SQL command can yield PGRES_COPY_BOTH, but + * defend against the future. PQexec() can't short-circuit + * it's way out of a PGRES_COPY_BOTH, so the connection will + * be useless at this point. XXX is there a method for + * clearing this status that's likely to work with every + * future command that can initiate it? + */ + psql_error("unexpected PQresultStatus (%d)", result_status); + return false; + + case PGRES_COPY_OUT: + case PGRES_COPY_IN: + is_copy = true; + break; + + case PGRES_EMPTY_QUERY: + case PGRES_COMMAND_OK: + case PGRES_TUPLES_OK: + is_copy = false; + break; + + default: + /* AcceptResult() should have caught anything else. */ + is_copy = false; + psql_error("unexpected PQresultStatus (%d)", result_status); + break; + } + + if (is_copy) + { + /* + * Marshal the COPY data. Either subroutine will get the + * connection out of its COPY state, then call PQresultStatus() + * once and report any error. + */ SetCancelConn(); - success = handleCopyOut(pset.db, pset.queryFout); + 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(); - break; - case PGRES_COPY_IN: - SetCancelConn(); - success = handleCopyIn(pset.db, pset.cur_cmd_source, - PQbinaryTuples(results)); - 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; + } - default: - break; - } + first_cycle = false; + } while (next_result); /* may need this to recover from conn loss during COPY */ if (!CheckConnection()) @@ -708,7 +787,7 @@ PrintQueryStatus(PGresult *results) static bool PrintQueryResults(PGresult *results) { - bool success = false; + bool success; const char *cmdstatus; if (!results) @@ -738,11 +817,21 @@ PrintQueryResults(PGresult *results) case PGRES_COPY_OUT: case PGRES_COPY_IN: + case PGRES_COPY_BOTH: /* nothing to do here */ success = true; break; + case PGRES_BAD_RESPONSE: + case PGRES_NONFATAL_ERROR: + case PGRES_FATAL_ERROR: + success = false; + break; + default: + success = false; + psql_error("unexpected PQresultStatus (%d)", + PQresultStatus(results)); break; } @@ -867,7 +956,7 @@ SendQuery(const char *query) /* these operations are included in the timing result: */ ResetCancelConn(); - OK = (AcceptResult(results) && ProcessCopyResult(results)); + OK = ProcessResult(&results); if (pset.timing) { @@ -877,7 +966,7 @@ SendQuery(const char *query) } /* but printing results isn't: */ - if (OK) + if (OK && results) OK = PrintQueryResults(results); } else @@ -891,34 +980,44 @@ SendQuery(const char *query) /* If we made a temporary savepoint, possibly release/rollback */ if (on_error_rollback_savepoint) { - const char *svptcmd; + const char *svptcmd = NULL; transaction_status = PQtransactionStatus(pset.db); - if (transaction_status == PQTRANS_INERROR) + switch (transaction_status) { - /* We always rollback on an error */ - svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint"; - } - else if (transaction_status != PQTRANS_INTRANS) - { - /* If they are no longer in a transaction, then do nothing */ - svptcmd = NULL; - } - else - { - /* - * Do nothing if they are messing with savepoints themselves: If - * the user did RELEASE or ROLLBACK, our savepoint is gone. If - * they issued a SAVEPOINT, releasing ours would remove theirs. - */ - if (results && - (strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 || - strcmp(PQcmdStatus(results), "RELEASE") == 0 || - strcmp(PQcmdStatus(results), "ROLLBACK") == 0)) - svptcmd = NULL; - else - svptcmd = "RELEASE pg_psql_temporary_savepoint"; + case PQTRANS_INERROR: + /* We always rollback on an error */ + svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint"; + break; + + case PQTRANS_IDLE: + /* If they are no longer in a transaction, then do nothing */ + break; + + case PQTRANS_INTRANS: + /* + * Do nothing if they are messing with savepoints themselves: + * If the user did RELEASE or ROLLBACK, our savepoint is + * gone. If they issued a SAVEPOINT, releasing ours would + * remove theirs. + */ + if (results && + (strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 || + strcmp(PQcmdStatus(results), "RELEASE") == 0 || + strcmp(PQcmdStatus(results), "ROLLBACK") == 0)) + svptcmd = NULL; + else + svptcmd = "RELEASE pg_psql_temporary_savepoint"; + break; + + case PQTRANS_ACTIVE: + case PQTRANS_UNKNOWN: + default: + OK = false; + psql_error("unexpected transaction status (%d)\n", + transaction_status); + break; } if (svptcmd) diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c index eaf633d2a5..a1dea9502c 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -244,8 +244,9 @@ do_copy(const char *args) { PQExpBufferData query; FILE *copystream; + FILE *save_file; + FILE **override_file; struct copy_options *options; - PGresult *result; bool success; struct stat st; @@ -261,6 +262,8 @@ 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) @@ -270,6 +273,8 @@ 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) @@ -308,52 +313,13 @@ do_copy(const char *args) if (options->after_tofrom) appendPQExpBufferStr(&query, options->after_tofrom); - result = PSQLexec(query.data, true); + /* 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); - switch (PQresultStatus(result)) - { - case PGRES_COPY_OUT: - SetCancelConn(); - success = handleCopyOut(pset.db, copystream); - ResetCancelConn(); - break; - case PGRES_COPY_IN: - SetCancelConn(); - success = handleCopyIn(pset.db, copystream, - PQbinaryTuples(result)); - ResetCancelConn(); - break; - case PGRES_NONFATAL_ERROR: - case PGRES_FATAL_ERROR: - case PGRES_BAD_RESPONSE: - success = false; - psql_error("\\copy: %s", PQerrorMessage(pset.db)); - break; - default: - success = false; - psql_error("\\copy: unexpected response (%d)\n", - PQresultStatus(result)); - break; - } - - PQclear(result); - - /* - * Make sure we have pumped libpq dry of results; else it may still be in - * ASYNC_BUSY state, leading to false readings in, eg, get_prompt(). - */ - while ((result = PQgetResult(pset.db)) != NULL) - { - success = false; - psql_error("\\copy: unexpected response (%d)\n", - PQresultStatus(result)); - /* if still in COPY IN state, try to get out of it */ - if (PQresultStatus(result) == PGRES_COPY_IN) - PQputCopyEnd(pset.db, _("trying to exit copy mode")); - PQclear(result); - } - if (options->file != NULL) { if (fclose(copystream) != 0) @@ -425,8 +391,30 @@ handleCopyOut(PGconn *conn, FILE *copystream) OK = false; } - /* Check command status and return to normal libpq state */ - res = PQgetResult(conn); + /* + * Check command status and return to normal libpq state. After a + * client-side error, the server will remain ready to deliver data. The + * cleanest thing is to fully drain and discard that data. If the + * client-side error happened early in a large file, this takes a long + * time. Instead, take advantage of the fact that PQexec() will silently + * end any ongoing PGRES_COPY_OUT state. This does cause us to lose the + * results of any commands following the COPY in a single command string. + * It also only works for protocol version 3. XXX should we clean up + * using the slow way when the connection is using protocol version 2? + * + * We must not ever return with the status still PGRES_COPY_OUT. Our + * caller is unable to distinguish that situation from reaching the next + * COPY in a command string that happened to contain two consecutive COPY + * TO STDOUT commands. We trust that no condition can make PQexec() fail + * indefinitely while retaining status PGRES_COPY_OUT. + */ + while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_OUT) + { + OK = false; + PQclear(res); + + PQexec(conn, "-- clear PGRES_COPY_OUT state"); + } if (PQresultStatus(res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); @@ -471,13 +459,8 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary) /* Terminate data transfer */ PQputCopyEnd(conn, _("canceled by user")); - /* Check command status and return to normal libpq state */ - res = PQgetResult(conn); - if (PQresultStatus(res) != PGRES_COMMAND_OK) - psql_error("%s", PQerrorMessage(conn)); - PQclear(res); - - return false; + OK = false; + goto copyin_cleanup; } /* Prompt if interactive input */ @@ -600,8 +583,24 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary) OK ? NULL : _("aborted because of read failure")) <= 0) OK = false; - /* Check command status and return to normal libpq state */ - res = PQgetResult(conn); +copyin_cleanup: + /* + * Check command status and return to normal libpq state + * + * We must not ever return with the status still PGRES_COPY_IN. Our + * caller is unable to distinguish that situation from reaching the next + * COPY in a command string that happened to contain two consecutive COPY + * FROM STDIN commands. XXX if something makes PQputCopyEnd() fail + * indefinitely while retaining status PGRES_COPY_IN, we get an infinite + * loop. This is more realistic than handleCopyOut()'s counterpart risk. + */ + while (res = PQgetResult(conn), PQresultStatus(res) == PGRES_COPY_IN) + { + OK = false; + PQclear(res); + + PQputCopyEnd(pset.db, _("trying to exit copy mode")); + } if (PQresultStatus(res) != PGRES_COMMAND_OK) { psql_error("%s", PQerrorMessage(conn)); diff --git a/src/test/regress/expected/copyselect.out b/src/test/regress/expected/copyselect.out index cbc140c538..7c16a61733 100644 --- a/src/test/regress/expected/copyselect.out +++ b/src/test/regress/expected/copyselect.out @@ -111,8 +111,6 @@ t \copy v_test1 to stdout ERROR: cannot copy from view "v_test1" HINT: Try the COPY (SELECT ...) TO variant. -\copy: ERROR: cannot copy from view "v_test1" -HINT: Try the COPY (SELECT ...) TO variant. -- -- Test \copy (select ...) -- @@ -124,3 +122,32 @@ HINT: Try the COPY (SELECT ...) TO variant. drop table test2; drop view v_test1; drop table test1; +-- psql handling of COPY in multi-command strings +copy (select 1) to stdout\; select 1/0; -- row, then error +1 +ERROR: division by zero +select 1/0\; copy (select 1) to stdout; -- error only +ERROR: division by zero +copy (select 1) to stdout\; copy (select 2) to stdout\; select 0\; select 3; -- 1 2 3 +1 +2 + ?column? +---------- + 3 +(1 row) + +create table test3 (c int); +select 0\; copy test3 from stdin\; copy test3 from stdin\; select 1; -- 1 + ?column? +---------- + 1 +(1 row) + +select * from test3; + c +--- + 1 + 2 +(2 rows) + +drop table test3; diff --git a/src/test/regress/sql/copyselect.sql b/src/test/regress/sql/copyselect.sql index 621d49444d..1d98dad3c8 100644 --- a/src/test/regress/sql/copyselect.sql +++ b/src/test/regress/sql/copyselect.sql @@ -80,3 +80,17 @@ copy (select t from test1 where id = 1) to stdout csv header force quote t; drop table test2; drop view v_test1; drop table test1; + +-- psql handling of COPY in multi-command strings +copy (select 1) to stdout\; select 1/0; -- row, then error +select 1/0\; copy (select 1) to stdout; -- error only +copy (select 1) to stdout\; copy (select 2) to stdout\; select 0\; select 3; -- 1 2 3 + +create table test3 (c int); +select 0\; copy test3 from stdin\; copy test3 from stdin\; select 1; -- 1 +1 +\. +2 +\. +select * from test3; +drop table test3;