diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 428f965c495..ea48f89e20b 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -1794,15 +1794,36 @@ ExecQueryAndProcessResults(const char *query, PsqlScanState scan_state; PQExpBuffer query_buf; promptStatus_t prompt_tmp; + PsqlScanResult scan_result; scan_state = psql_scan_create(&psqlscan_callbacks); psql_scan_setup(scan_state, query, strlen(query), pset.encoding, standard_strings()); query_buf = createPQExpBuffer(); - (void) psql_scan(scan_state, query_buf, &prompt_tmp); + /* + * A semicolon ends only one sub-command; keep scanning so that COPY + * FROM STDIN commands past the first semicolon are counted too. The + * count accumulates in scan_state across the psql_scan() calls. + */ + do + { + scan_result = psql_scan(scan_state, query_buf, &prompt_tmp); + } while (scan_result == PSCAN_SEMICOLON); - num_copy_from_stdin = psql_scan_count_copy_from_stdin(scan_state); + /* + * We expect the result now to be PSCAN_EOL. If it is PSCAN_BACKSLASH + * or PSCAN_INCOMPLETE, the server will get a parse error and refuse + * to execute any part of the command string, so don't expect any + * PGRES_COPY_IN results. (This will mean that we don't attempt to + * discard any following data, but this seems consistent with the + * general contract of psql_scan_count_copy_from_stdin, which is that + * it only promises to count syntactically-valid COPY commands.) + */ + if (scan_result == PSCAN_EOL) + num_copy_from_stdin = psql_scan_count_copy_from_stdin(scan_state); + else + num_copy_from_stdin = 0; destroyPQExpBuffer(query_buf); psql_scan_destroy(scan_state); diff --git a/src/bin/psql/t/001_basic.pl b/src/bin/psql/t/001_basic.pl index a335acca68c..46798f490c8 100644 --- a/src/bin/psql/t/001_basic.pl +++ b/src/bin/psql/t/001_basic.pl @@ -532,6 +532,36 @@ psql_fails_like( qr/COPY in a pipeline is not supported, aborting connection/, '\copy to in pipeline: fails'); +# Test execution of COPY FROM STDIN in -c. This case is a bit weird +# because it will read from psql's stdin not from the command source. +# To make it even weirder, try two such commands, to stress psql's logic +# that counts them. Also test both \. and EOF termination. +{ + $node->safe_psql('postgres', 'CREATE TABLE copy_stdin_count (a int)'); + my ($stdin, $stdout, $stderr) = ("50\n\\.\n60\n", '', ''); + my $ret = IPC::Run::run( + [ + 'psql', '--no-psqlrc', + '--set' => 'ON_ERROR_STOP=1', + '--dbname' => $node->connstr('postgres'), + '--command' => + 'COPY copy_stdin_count FROM STDIN; COPY copy_stdin_count FROM STDIN', + ], + '<' => \$stdin, + '>' => \$stdout, + '2>' => \$stderr); + + ok($ret, '-c COPY FROM STDIN: psql exits 0'); + unlike( + $stderr, + qr/unexpected COPY_IN result/, + '-c COPY FROM STDIN: unexpected COPY_IN result'); + + my $data = $node->safe_psql('postgres', 'SELECT * FROM copy_stdin_count'); + is($data, "50\n60", '-c COPY FROM STDIN: correct data loaded'); +} + +# Test \restrict and \unrestrict. psql_fails_like( $node, qq{\\restrict test