diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 9d9d04b8265..1a097772b32 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -27,6 +27,7 @@ #include "fe_utils/cancel.h" #include "fe_utils/mbprint.h" #include "fe_utils/string_utils.h" +#include "mainloop.h" #include "portability/instr_time.h" #include "settings.h" @@ -35,6 +36,7 @@ static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec); static int ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_gone_p, + int num_copy_from_stdin, bool is_watch, const printQueryOpt *opt, FILE *printQueryFout); @@ -626,7 +628,11 @@ PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout) SetCancelConn(pset.db); - res = ExecQueryAndProcessResults(query, &elapsed_msec, NULL, true, opt, printQueryFout); + res = ExecQueryAndProcessResults(query, + &elapsed_msec, NULL, + -1, + true, + opt, printQueryFout); ResetCancelConn(); @@ -801,7 +807,7 @@ ExecQueryTuples(const PGresult *result) fflush(stdout); } - if (!SendQuery(query)) + if (!SendQuery(query, -1)) { /* Error - abandon execution if ON_ERROR_STOP */ success = false; @@ -832,8 +838,7 @@ loop_exit: * once and report any error. Return whether all was ok. * * For COPY OUT, direct the output to copystream, or discard if that's NULL. - * For COPY IN, use pset.copyStream as data source if it's set, - * otherwise cur_cmd_source. + * For COPY IN, read from copystream (which mustn't be NULL). * * Update *resultp if further processing is necessary; set to NULL otherwise. * Return a result when queryFout can safely output a result status: on COPY @@ -874,8 +879,7 @@ HandleCopyResult(PGresult **resultp, FILE *copystream) else { /* COPY IN */ - /* Ignore the copystream argument passed to the function */ - copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source; + Assert(copystream); success = handleCopyIn(pset.db, copystream, PQbinaryTuples(*resultp), @@ -1016,10 +1020,14 @@ PrintQueryResult(PGresult *result, bool last, * To send "back door" queries (generated by slash commands, etc.) in a * controlled way, use PSQLexec(). * + * In the most common case, the caller can determine whether the query + * includes any COPY FROM STDIN command(s); if so, pass the count of them. + * Otherwise pass num_copy_from_stdin == -1 and we'll compute it locally. + * * Returns true if the query executed successfully, false otherwise. */ bool -SendQuery(const char *query) +SendQuery(const char *query, int num_copy_from_stdin) { bool timing = pset.timing; PGTransactionStatusType transaction_status; @@ -1114,7 +1122,11 @@ SendQuery(const char *query) pset.crosstab_flag || !is_select_command(query)) { /* Default fetch-it-all-and-print mode */ - OK = (ExecQueryAndProcessResults(query, &elapsed_msec, &svpt_gone, false, NULL, NULL) > 0); + OK = (ExecQueryAndProcessResults(query, + &elapsed_msec, &svpt_gone, + num_copy_from_stdin, + false, + NULL, NULL) > 0); } else { @@ -1383,6 +1395,7 @@ DescribeQuery(const char *query, double *elapsed_msec) static int ExecQueryAndProcessResults(const char *query, double *elapsed_msec, bool *svpt_gone_p, + int num_copy_from_stdin, bool is_watch, const printQueryOpt *opt, FILE *printQueryFout) { @@ -1422,6 +1435,29 @@ ExecQueryAndProcessResults(const char *query, return 0; } + /* + * If the caller didn't count the number of COPY FROM STDIN command(s) in + * the query string, we must do so now. + */ + if (num_copy_from_stdin < 0) + { + PsqlScanState scan_state; + PQExpBuffer query_buf; + promptStatus_t prompt_tmp; + + 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); + + num_copy_from_stdin = psql_scan_count_copy_from_stdin(scan_state); + + destroyPQExpBuffer(query_buf); + psql_scan_destroy(scan_state); + } + /* first result */ result = PQgetResult(pset.db); @@ -1505,14 +1541,15 @@ ExecQueryAndProcessResults(const char *query, { FILE *copy_stream = NULL; - /* - * For COPY OUT, direct the output to the default place (probably - * a pager pipe) for \watch, or to pset.copyStream for \copy, - * otherwise to pset.gfname if that's set, otherwise to - * pset.queryFout. - */ + /* Select data source or sink. */ if (result_status == PGRES_COPY_OUT) { + /* + * For COPY OUT, direct the output to the default place + * (probably a pager pipe) for \watch, or to pset.copyStream + * for \copy, otherwise to pset.gfname if that's set, + * otherwise to pset.queryFout. + */ if (is_watch) { /* invoked by \watch */ @@ -1547,6 +1584,32 @@ ExecQueryAndProcessResults(const char *query, copy_stream = pset.queryFout; } } + else + { + if (num_copy_from_stdin <= 0) + { + /* + * We didn't send a COPY FROM STDIN command. The server + * is broken or possibly malicious. Report and quit. (We + * could instead send an empty COPY response, but it's not + * clear that that's a better behavior; it would make it + * harder to diagnose any such problem.) + */ + pg_log_info("unexpected COPY_IN result, aborting connection"); + exit(EXIT_BADCONN); + } + /* Keep track of the number of remaining copy operations */ + num_copy_from_stdin--; + + /* + * For COPY IN, read from pset.copyStream if \copy set that, + * otherwise from the current command source. + */ + if (pset.copyStream) + copy_stream = pset.copyStream; + else + copy_stream = pset.cur_cmd_source; + } /* * Even if the output stream could not be opened, we call @@ -1648,6 +1711,24 @@ ExecQueryAndProcessResults(const char *query, if (!CheckConnection()) return -1; + /* + * If we were expecting COPY ... FROM STDIN, and we didn't get + * PGRES_COPY_IN (presumably because the COPY failed server-side), and the + * data is supposed to come from the current command source, we must eat + * up the data in order to stay in sync with the command file's contents. + * Repeat for the number of unfulfilled COPY commands. + */ + while (num_copy_from_stdin > 0) + { + if (pset.copyStream == NULL || + pset.copyStream == pset.cur_cmd_source) + (void) handleCopyIn(NULL, /* no connection */ + pset.cur_cmd_source, + false, /* XXX assume not binary */ + NULL); /* no result */ + num_copy_from_stdin--; + } + return cancel_pressed ? 0 : success ? 1 : -1; } diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h index d84c3a007cf..ccb1384693d 100644 --- a/src/bin/psql/common.h +++ b/src/bin/psql/common.h @@ -31,7 +31,7 @@ extern void psql_setup_cancel_handler(void); extern PGresult *PSQLexec(const char *query); extern int PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout); -extern bool SendQuery(const char *query); +extern bool SendQuery(const char *query, int num_copy_from_stdin); extern bool is_superuser(void); extern bool standard_strings(void); diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c index 5db516c4f89..cae0bd1cc72 100644 --- a/src/bin/psql/copy.c +++ b/src/bin/psql/copy.c @@ -369,7 +369,7 @@ do_copy(const char *args) /* run it like a user command, but with copystream as data source/sink */ pset.copyStream = copystream; - success = SendQuery(query.data); + success = SendQuery(query.data, options->from ? 1 : 0); pset.copyStream = NULL; termPQExpBuffer(&query); @@ -497,11 +497,13 @@ handleCopyOut(PGconn *conn, FILE *copystream, PGresult **res) * sends data to complete a COPY ... FROM STDIN command * * conn should be a database connection that you just issued COPY FROM on - * and got back a PGRES_COPY_IN result. + * and got back a PGRES_COPY_IN result. Alternatively, if conn is NULL, + * we read and discard the appropriate amount of data from copystream. * copystream is the file stream to read the data from. * isbinary can be set from PQbinaryTuples(). - * The final status for the COPY is returned into *res (but note - * we already reported the error, if it's not a success result). + * The final status for the COPY is returned into *res; but note + * we already reported the error, if it's not a success result. + * Also, if conn is NULL then *res is not touched. * * result is true if successful, false if not. */ @@ -516,6 +518,12 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) char buf[COPYBUFSIZ]; bool showprompt; + /* We want to prompt if interactive input ... */ + showprompt = isatty(fileno(copystream)); + /* ... but if we're just discarding data, don't bother the user at all */ + if (showprompt && !conn) + return true; + /* * Establish longjmp destination for exiting from wait-for-input. (This is * only effective while sigint_interrupt_enabled is TRUE.) @@ -525,24 +533,19 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) /* got here with longjmp */ /* Terminate data transfer */ - PQputCopyEnd(conn, - (PQprotocolVersion(conn) < 3) ? NULL : - _("canceled by user")); + if (conn) + PQputCopyEnd(conn, + (PQprotocolVersion(conn) < 3) ? NULL : + _("canceled by user")); OK = false; goto copyin_cleanup; } - /* Prompt if interactive input */ - if (isatty(fileno(copystream))) - { - showprompt = true; - if (!pset.quiet) - puts(_("Enter data to be copied followed by a newline.\n" - "End with a backslash and a period on a line by itself, or an EOF signal.")); - } - else - showprompt = false; + /* Issue initial prompt if interactive input */ + if (showprompt && !pset.quiet) + puts(_("Enter data to be copied followed by a newline.\n" + "End with a backslash and a period on a line by itself, or an EOF signal.")); OK = true; @@ -571,7 +574,7 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) if (buflen <= 0) break; - if (PQputCopyData(conn, buf, buflen) <= 0) + if (conn && PQputCopyData(conn, buf, buflen) <= 0) { OK = false; break; @@ -659,7 +662,7 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) */ if (buflen >= COPYBUFSIZ - 5 || (copydone && buflen > 0)) { - if (PQputCopyData(conn, buf, buflen) <= 0) + if (conn && PQputCopyData(conn, buf, buflen) <= 0) { OK = false; break; @@ -680,7 +683,8 @@ handleCopyIn(PGconn *conn, FILE *copystream, bool isbinary, PGresult **res) * keep the version checks just in case you're using a pre-v14 libpq.so at * runtime) */ - if (PQputCopyEnd(conn, + if (conn && + PQputCopyEnd(conn, (OK || PQprotocolVersion(conn) < 3) ? NULL : _("aborted because of read failure")) <= 0) OK = false; @@ -697,6 +701,10 @@ copyin_cleanup: */ clearerr(copystream); + /* Done if we don't have a connection to clean up */ + if (!conn) + return OK; + /* * Check command status and return to normal libpq state. * diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c index b0c4177a20e..58e9ab4a892 100644 --- a/src/bin/psql/mainloop.c +++ b/src/bin/psql/mainloop.c @@ -436,7 +436,9 @@ MainLoop(FILE *source) /* execute query unless we're in an inactive \if branch */ if (conditional_active(cond_stack)) { - success = SendQuery(query_buf->data); + /* count_copy_from_stdin should be reliable here */ + success = SendQuery(query_buf->data, + psql_scan_count_copy_from_stdin(scan_state)); slashCmdStatus = success ? PSQL_CMD_SEND : PSQL_CMD_ERROR; pset.stmt_lineno = 1; @@ -448,9 +450,10 @@ MainLoop(FILE *source) query_buf = swap_buf; } resetPQExpBuffer(query_buf); + /* reset parsing state, too */ + psql_scan_reset(scan_state); added_nl_pos = -1; - /* we need not do psql_scan_reset() here */ } else { @@ -512,7 +515,7 @@ MainLoop(FILE *source) /* should not see this in inactive branch */ Assert(conditional_active(cond_stack)); - success = SendQuery(query_buf->data); + success = SendQuery(query_buf->data, -1); /* transfer query to previous_buf by pointer-swapping */ { @@ -522,8 +525,7 @@ MainLoop(FILE *source) query_buf = swap_buf; } resetPQExpBuffer(query_buf); - - /* flush any paren nesting info after forced send */ + /* reset parsing state, too */ psql_scan_reset(scan_state); } else if (slashCmdStatus == PSQL_CMD_NEWEDIT) @@ -610,7 +612,8 @@ MainLoop(FILE *source) /* execute query unless we're in an inactive \if branch */ if (conditional_active(cond_stack)) { - success = SendQuery(query_buf->data); + success = SendQuery(query_buf->data, + psql_scan_count_copy_from_stdin(scan_state)); } else { diff --git a/src/bin/psql/psqlscanslash.l b/src/bin/psql/psqlscanslash.l index c85c6fdbd47..a917d1ec180 100644 --- a/src/bin/psql/psqlscanslash.l +++ b/src/bin/psql/psqlscanslash.l @@ -703,14 +703,15 @@ PsqlScanStateSave * psql_scan_get_lex_state(PsqlScanState state) { PsqlScanStateSave *lex_state = pg_malloc_object(PsqlScanStateSave); - StaticAssertStmt(sizeof(lex_state->identifiers) == sizeof(state->identifiers), - "identifiers array lengths must match"); + StaticAssertStmt(sizeof(lex_state->init_idents) == sizeof(state->init_idents), + "init_idents array lengths must match"); lex_state->paren_depth = state->paren_depth; lex_state->begin_depth = state->begin_depth; - lex_state->identifier_count = state->identifier_count; - memcpy(lex_state->identifiers, state->identifiers, - sizeof(lex_state->identifiers)); + lex_state->copy_stdin_count = state->copy_stdin_count; + lex_state->init_idents_count = state->init_idents_count; + memcpy(lex_state->init_idents, state->init_idents, + sizeof(lex_state->init_idents)); return lex_state; } @@ -723,9 +724,10 @@ psql_scan_set_lex_state(PsqlScanState state, { state->paren_depth = lex_state->paren_depth; state->begin_depth = lex_state->begin_depth; - state->identifier_count = lex_state->identifier_count; - memcpy(state->identifiers, lex_state->identifiers, - sizeof(state->identifiers)); + state->copy_stdin_count = lex_state->copy_stdin_count; + state->init_idents_count = lex_state->init_idents_count; + memcpy(state->init_idents, lex_state->init_idents, + sizeof(state->init_idents)); } /* diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c index 7c2f555f15c..8fb73c43269 100644 --- a/src/bin/psql/startup.c +++ b/src/bin/psql/startup.c @@ -381,7 +381,7 @@ main(int argc, char *argv[]) if (pset.echo == PSQL_ECHO_ALL) puts(cell->val); - successResult = SendQuery(cell->val) + successResult = SendQuery(cell->val, -1) ? EXIT_SUCCESS : EXIT_FAILURE; } else if (cell->action == ACT_SINGLE_SLASH) diff --git a/src/fe_utils/psqlscan.l b/src/fe_utils/psqlscan.l index 6c603f6e7a0..2ce1bff890f 100644 --- a/src/fe_utils/psqlscan.l +++ b/src/fe_utils/psqlscan.l @@ -76,6 +76,10 @@ typedef int YYSTYPE; extern int psql_yyget_column(yyscan_t yyscanner); extern void psql_yyset_column(int column_no, yyscan_t yyscanner); +static bool psqlscan_is_copy_from_stdin(PsqlScanState state); +static void psqlscan_track_identifier(PsqlScanState state, + const char *identifier); + %} %option reentrant @@ -650,11 +654,15 @@ other . ";" { ECHO; - if (cur_state->paren_depth == 0 && cur_state->begin_depth == 0) + if (cur_state->paren_depth == 0 && + cur_state->begin_depth == 0) { + /* Remember if this subcommand was COPY FROM STDIN */ + if (psqlscan_is_copy_from_stdin(cur_state)) + cur_state->copy_stdin_count++; /* Terminate lexing temporarily */ cur_state->start_state = YY_START; - cur_state->identifier_count = 0; + cur_state->init_idents_count = 0; return LEXRES_SEMI; } } @@ -667,8 +675,16 @@ other . "\\"[;:] { /* Force a semi-colon or colon into the query buffer */ psqlscan_emit(cur_state, yytext + 1, 1); - if (yytext[1] == ';') - cur_state->identifier_count = 0; + /* Reset BEGIN/END/COPY tracking if semi at outer level */ + if (yytext[1] == ';' && + cur_state->paren_depth == 0 && + cur_state->begin_depth == 0) + { + /* Remember if this subcommand was COPY FROM STDIN */ + if (psqlscan_is_copy_from_stdin(cur_state)) + cur_state->copy_stdin_count++; + cur_state->init_idents_count = 0; + } } "\\" { @@ -876,55 +892,7 @@ other . {identifier} { - /* - * We need to track if we are inside a BEGIN .. END block - * in a function definition, so that semicolons contained - * therein don't terminate the whole statement. Short of - * writing a full parser here, the following heuristic - * should work. First, we track whether the beginning of - * the statement matches CREATE [OR REPLACE] - * {FUNCTION|PROCEDURE} - */ - - if (cur_state->identifier_count == 0) - memset(cur_state->identifiers, 0, sizeof(cur_state->identifiers)); - - if (pg_strcasecmp(yytext, "create") == 0 || - pg_strcasecmp(yytext, "function") == 0 || - pg_strcasecmp(yytext, "procedure") == 0 || - pg_strcasecmp(yytext, "or") == 0 || - pg_strcasecmp(yytext, "replace") == 0) - { - if (cur_state->identifier_count < sizeof(cur_state->identifiers)) - cur_state->identifiers[cur_state->identifier_count] = pg_tolower((unsigned char) yytext[0]); - } - - cur_state->identifier_count++; - - if (cur_state->identifiers[0] == 'c' && - (cur_state->identifiers[1] == 'f' || cur_state->identifiers[1] == 'p' || - (cur_state->identifiers[1] == 'o' && cur_state->identifiers[2] == 'r' && - (cur_state->identifiers[3] == 'f' || cur_state->identifiers[3] == 'p'))) && - cur_state->paren_depth == 0) - { - if (pg_strcasecmp(yytext, "begin") == 0) - cur_state->begin_depth++; - else if (pg_strcasecmp(yytext, "case") == 0) - { - /* - * CASE also ends with END. We only need to track - * this if we are already inside a BEGIN. - */ - if (cur_state->begin_depth >= 1) - cur_state->begin_depth++; - } - else if (pg_strcasecmp(yytext, "end") == 0) - { - if (cur_state->begin_depth > 0) - cur_state->begin_depth--; - } - } - + psqlscan_track_identifier(cur_state, yytext); ECHO; } @@ -951,6 +919,145 @@ other . /* LCOV_EXCL_STOP */ +/* + * Record the first few keywords/identifiers of a statement + * in the idents[] array, of length idents_size. + * *idents_count is the number of entries filled so far. + * + * We record the interesting keywords using their first character, which + * works so long as those are all different. We could switch to an enum + * if that stops being true, but for now this is easy and compact. + */ +static void +psqlscan_record_initial_keyword(const char *identifier, + char *idents, + int idents_size, + int *idents_count) +{ + if (*idents_count < idents_size) + { + /* + * What we need to recognize is CREATE [OR REPLACE] FUNCTION/PROCEDURE. + * We record these keywords in lower case. We also need to recognize + * COPY ... FROM STDIN. (Note: the backend grammar doesn't + * distinguish STDIN from STDOUT, so we should not do so here either.) + * We record these keywords in upper case, to avoid conflicting with + * the first set. + */ + if (pg_strcasecmp(identifier, "create") == 0 || + pg_strcasecmp(identifier, "function") == 0 || + pg_strcasecmp(identifier, "procedure") == 0 || + pg_strcasecmp(identifier, "or") == 0 || + pg_strcasecmp(identifier, "replace") == 0) + idents[*idents_count] = pg_tolower((unsigned char) identifier[0]); + else if (pg_strcasecmp(identifier, "copy") == 0 || + pg_strcasecmp(identifier, "from") == 0 || + pg_strcasecmp(identifier, "stdin") == 0 || + pg_strcasecmp(identifier, "stdout") == 0) + idents[*idents_count] = pg_toupper((unsigned char) identifier[0]); + /* For other keywords or identifiers, leave '\0' in the array entry */ + (*idents_count)++; + } +} + +/* + * Does the current input match CREATE [OR REPLACE] {FUNCTION|PROCEDURE}? + */ +static bool +psqlscan_is_create_routine(const char *idents) +{ + return idents[0] == 'c' && + (idents[1] == 'f' || idents[1] == 'p' || + (idents[1] == 'o' && idents[2] == 'r' && + (idents[3] == 'f' || idents[3] == 'p'))); +} + +/* + * Does the current input match COPY ... FROM STDIN? + */ +static bool +psqlscan_is_copy_from_stdin(PsqlScanState state) +{ + const char *idents = state->init_idents; + + /* + * The first word must be COPY, but after that there could be up to four + * identifiers (BINARY database.schema.table) before FROM. Since some of + * the words we track are not reserved words, don't assume the intervening + * array entries are '\0'. Life is simplified here by the fact that + * psqlscan_track_identifier ignores everything within parens: we won't + * see column lists nor the query in COPY (query). + */ + if (idents[0] != 'C') + return false; + for (int i = 1; i < lengthof(state->init_idents) - 1; i++) + { + /* Scan to find FROM; if not seen within range, it's not valid COPY */ + if (idents[i] != 'F') + continue; + /* It's COPY FROM STDIN only if the next word is STDIN */ + return (idents[i + 1] == 'S'); + } + return false; +} + +/* + * This function is called each time the lexer recognizes an unquoted + * identifier (which could also be a keyword, and indeed keywords are the + * only case we really care about here). It presently has two tasks: + * + * 1. Track whether we are inside BEGIN .. END in a function definition, + * so that semicolons contained therein don't terminate the whole statement. + * Short of writing a full parser here, the following heuristic should work. + * + * We track whether the beginning of the statement matches CREATE [OR REPLACE] + * {FUNCTION|PROCEDURE}. If so, count BEGIN and END pairs. We also have to + * account for CASE ... END. + * + * 2. Record enough information for psqlscan_is_copy_from_stdin() to recognize + * COPY FROM STDIN commands. + */ +static void +psqlscan_track_identifier(PsqlScanState state, const char *identifier) +{ + /* None of this needs to happen when we're inside parentheses */ + if (state->paren_depth != 0) + return; + + /* Reset all my state at the start of each new statement */ + if (state->init_idents_count == 0) + memset(state->init_idents, 0, sizeof(state->init_idents)); + + /* Record initial keywords if init_idents_count is small enough */ + psqlscan_record_initial_keyword(identifier, + state->init_idents, + lengthof(state->init_idents), + &state->init_idents_count); + + /* + * Track BEGIN/CASE/END only when within an appropriate statement. + */ + if (psqlscan_is_create_routine(state->init_idents)) + { + if (pg_strcasecmp(identifier, "begin") == 0) + state->begin_depth++; + else if (pg_strcasecmp(identifier, "case") == 0) + { + /* + * CASE also ends with END. We only need to track this if we are + * already inside a BEGIN. + */ + if (state->begin_depth >= 1) + state->begin_depth++; + } + else if (pg_strcasecmp(identifier, "end") == 0) + { + if (state->begin_depth > 0) + state->begin_depth--; + } + } +} + /* * Create a lexer working state struct. * @@ -1216,10 +1323,9 @@ psql_scan_finish(PsqlScanState state) /* * Reset lexer scanning state to start conditions. This is appropriate * for executing \r psql commands (or any other time that we discard the - * prior contents of query_buf). It is not, however, necessary to do this - * when we execute and clear the buffer after getting a PSCAN_SEMICOLON or - * PSCAN_EOL scan result, because the scan state must be INITIAL when those - * conditions are returned. + * prior contents of query_buf). Do not call this between psql_scan() + * calls that are scanning successive chunks of a single query string; + * do call it when preparing to process a new query string. * * Note that this is unrelated to flushing unread input; that task is * done by psql_scan_finish(). @@ -1233,8 +1339,9 @@ psql_scan_reset(PsqlScanState state) if (state->dolqstart) free(state->dolqstart); state->dolqstart = NULL; - state->identifier_count = 0; state->begin_depth = 0; + state->copy_stdin_count = 0; + state->init_idents_count = 0; } /* @@ -1258,6 +1365,26 @@ psql_scan_reselect_sql_lexer(PsqlScanState state) state->start_state = INITIAL; } +/* + * Return the number of COPY ... FROM STDIN commands in the input string. + * + * This should be called only after we've finished parsing a complete + * string and are ready to send it to the backend. + */ +int +psql_scan_count_copy_from_stdin(PsqlScanState state) +{ + if (state->init_idents_count > 0) + { + /* Count any COPY FROM STDIN following the last semicolon */ + if (psqlscan_is_copy_from_stdin(state)) + state->copy_stdin_count++; + /* ... but do so only once */ + state->init_idents_count = 0; + } + return state->copy_stdin_count; +} + /* * Return true if lexer is currently in an "inside quotes" state. * diff --git a/src/include/fe_utils/psqlscan.h b/src/include/fe_utils/psqlscan.h index 8155a369879..20dd08141a9 100644 --- a/src/include/fe_utils/psqlscan.h +++ b/src/include/fe_utils/psqlscan.h @@ -88,6 +88,8 @@ extern void psql_scan_reset(PsqlScanState state); extern void psql_scan_reselect_sql_lexer(PsqlScanState state); +extern int psql_scan_count_copy_from_stdin(PsqlScanState state); + extern bool psql_scan_in_quote(PsqlScanState state); #endif /* PSQLSCAN_H */ diff --git a/src/include/fe_utils/psqlscan_int.h b/src/include/fe_utils/psqlscan_int.h index 4bde8db443c..08385c6d788 100644 --- a/src/include/fe_utils/psqlscan_int.h +++ b/src/include/fe_utils/psqlscan_int.h @@ -116,12 +116,14 @@ typedef struct PsqlScanStateData char *dolqstart; /* current $foo$ quote start string */ /* - * State to track boundaries of BEGIN ... END blocks in function - * definitions, so that semicolons do not send query too early. + * State used to track boundaries of BEGIN ... END blocks in function + * definitions, so that semicolons do not send query too early. We also + * use this state to detect and count COPY FROM STDIN commands. */ - int identifier_count; /* identifiers since start of statement */ - char identifiers[4]; /* records the first few identifiers */ int begin_depth; /* depth of begin/end pairs */ + int copy_stdin_count; /* number of COPY FROM STDIN commands */ + int init_idents_count; /* # identifiers since start of statement */ + char init_idents[8]; /* records the first few identifiers */ /* * Callback functions provided by the program making use of the lexer, @@ -144,8 +146,9 @@ struct PsqlScanStateSave { int paren_depth; /* depth of nesting in parentheses */ int begin_depth; /* depth of begin/end pairs */ - int identifier_count; /* identifiers since start of statement */ - char identifiers[4]; /* records the first few identifiers */ + int copy_stdin_count; /* number of COPY FROM STDIN commands */ + int init_idents_count; /* # identifiers since start of statement */ + char init_idents[8]; /* records the first few identifiers */ }; diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out index 8c5d6c2f93a..cf3a7a0c29d 100644 --- a/src/test/regress/expected/psql.out +++ b/src/test/regress/expected/psql.out @@ -5517,6 +5517,11 @@ CONTEXT: PL/pgSQL function warn(text) line 2 at RAISE \set SHOW_ALL_RESULTS on DROP FUNCTION warn(TEXT); +-- \copy must skip in-line data, even if the issued COPY command fails. +\copy no_such_table from stdin +ERROR: relation "no_such_table" does not exist +\echo this should get output +this should get output -- -- AUTOCOMMIT and combined queries -- diff --git a/src/test/regress/sql/alter_table.sql b/src/test/regress/sql/alter_table.sql index 9894b159a1a..2bbabd4b528 100644 --- a/src/test/regress/sql/alter_table.sql +++ b/src/test/regress/sql/alter_table.sql @@ -1146,7 +1146,9 @@ copy attest from stdin; \. select * from attest; copy attest(a) from stdin; +\. copy attest("........pg.dropped.1........") from stdin; +\. copy attest(b,c) from stdin; 31 32 \. diff --git a/src/test/regress/sql/copy.sql b/src/test/regress/sql/copy.sql index f9da7b1508f..e773202222b 100644 --- a/src/test/regress/sql/copy.sql +++ b/src/test/regress/sql/copy.sql @@ -209,6 +209,7 @@ alter table header_copytest drop column c; alter table header_copytest add column c text; copy header_copytest to stdout with (header match); copy header_copytest from stdin with (header wrong_choice); +\. -- works copy header_copytest from stdin with (header match); a b c diff --git a/src/test/regress/sql/copy2.sql b/src/test/regress/sql/copy2.sql index b3c16af48ee..f34c9e4c95e 100644 --- a/src/test/regress/sql/copy2.sql +++ b/src/test/regress/sql/copy2.sql @@ -52,23 +52,37 @@ COPY x (a, b, c, d, e) from stdin; -- non-existent column in column list: should fail COPY x (xyz) from stdin; +\. -- redundant options COPY x from stdin (format CSV, FORMAT CSV); +\. COPY x from stdin (freeze off, freeze on); +\. COPY x from stdin (delimiter ',', delimiter ','); +\. COPY x from stdin (null ' ', null ' '); +\. COPY x from stdin (header off, header on); +\. COPY x from stdin (quote ':', quote ':'); +\. COPY x from stdin (escape ':', escape ':'); +\. COPY x from stdin (force_quote (a), force_quote *); +\. COPY x from stdin (force_not_null (a), force_not_null (b)); +\. COPY x from stdin (force_null (a), force_null (b)); +\. COPY x from stdin (convert_selectively (a), convert_selectively (b)); +\. COPY x from stdin (encoding 'sql_ascii', encoding 'sql_ascii'); +\. -- too many columns in column list: should fail COPY x (a, b, c, d, e, d, c) from stdin; +\. -- missing data: should fail COPY x from stdin; @@ -125,14 +139,19 @@ COPY x from stdin WHERE a > 60003; \. COPY x from stdin WHERE f > 60003; +\. COPY x from stdin WHERE a = max(x.b); +\. COPY x from stdin WHERE a IN (SELECT 1 FROM x); +\. COPY x from stdin WHERE a IN (generate_series(1,5)); +\. COPY x from stdin WHERE a = row_number() over(b); +\. -- check results of copy in @@ -329,10 +348,12 @@ ROLLBACK; -- should fail with "not referenced by COPY" error BEGIN; COPY forcetest (d, e) FROM STDIN WITH (FORMAT csv, FORCE_NOT_NULL(b)); +\. ROLLBACK; -- should fail with "not referenced by COPY" error BEGIN; COPY forcetest (d, e) FROM STDIN WITH (FORMAT csv, FORCE_NULL(b)); +\. ROLLBACK; \pset null '' diff --git a/src/test/regress/sql/copyselect.sql b/src/test/regress/sql/copyselect.sql index e32a4f8e38e..273a5f90a56 100644 --- a/src/test/regress/sql/copyselect.sql +++ b/src/test/regress/sql/copyselect.sql @@ -42,6 +42,7 @@ copy (select t into temp test3 from test1 where id=3) to stdout; -- This should fail -- copy (select * from test1) from stdin; +\. -- -- This should fail -- diff --git a/src/test/regress/sql/generated.sql b/src/test/regress/sql/generated.sql index ec728f56d34..1b45348ad60 100644 --- a/src/test/regress/sql/generated.sql +++ b/src/test/regress/sql/generated.sql @@ -203,10 +203,13 @@ COPY gtest1 FROM stdin; \. COPY gtest1 (a, b) FROM stdin; +\. COPY gtest1 FROM stdin WHERE b <> 10; +\. COPY gtest1 FROM stdin WHERE gtest1 IS NULL; +\. SELECT * FROM gtest1 ORDER BY a; @@ -223,6 +226,7 @@ COPY gtest3 FROM stdin; \. COPY gtest3 (a, b) FROM stdin; +\. SELECT * FROM gtest3 ORDER BY a; diff --git a/src/test/regress/sql/privileges.sql b/src/test/regress/sql/privileges.sql index 8aa62ad23e3..ea3ddf287c9 100644 --- a/src/test/regress/sql/privileges.sql +++ b/src/test/regress/sql/privileges.sql @@ -172,6 +172,7 @@ BEGIN; LOCK atest2 IN ACCESS EXCLUSIVE MODE; -- fail COMMIT; COPY atest2 FROM stdin; -- fail +\. GRANT ALL ON atest1 TO PUBLIC; -- fail -- checks in subquery, both ok @@ -215,6 +216,7 @@ BEGIN; LOCK atest2 IN ACCESS EXCLUSIVE MODE; -- ok COMMIT; COPY atest2 FROM stdin; -- fail +\. -- checks in subquery, both fail SELECT * FROM atest1 WHERE ( b IN ( SELECT col1 FROM atest2 ) ); @@ -447,6 +449,7 @@ SELECT one, two FROM atest5 NATURAL JOIN atest6; -- ok now -- test column-level privileges for INSERT and UPDATE INSERT INTO atest5 (two) VALUES (3); -- ok COPY atest5 FROM stdin; -- fail +\. COPY atest5 (two) FROM stdin; -- ok 1 \. diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql index e743d364bff..ee6d9882534 100644 --- a/src/test/regress/sql/psql.sql +++ b/src/test/regress/sql/psql.sql @@ -1405,6 +1405,15 @@ SELECT 1 AS one \; SELECT warn('1.5') \; SELECT 2 AS two ; \set SHOW_ALL_RESULTS on DROP FUNCTION warn(TEXT); +-- \copy must skip in-line data, even if the issued COPY command fails. +\copy no_such_table from stdin +foo +\echo this should not get output +bar +\echo this should not get output +\. +\echo this should get output + -- -- AUTOCOMMIT and combined queries -- diff --git a/src/test/regress/sql/rowsecurity.sql b/src/test/regress/sql/rowsecurity.sql index 16405991b6e..bb4aae3527b 100644 --- a/src/test/regress/sql/rowsecurity.sql +++ b/src/test/regress/sql/rowsecurity.sql @@ -1632,8 +1632,10 @@ COPY copy_t FROM STDIN; --ok SET SESSION AUTHORIZATION regress_rls_bob; SET row_security TO OFF; COPY copy_t FROM STDIN; --fail - would be affected by RLS. +\. SET row_security TO ON; COPY copy_t FROM STDIN; --fail - COPY FROM not supported by RLS. +\. -- Check COPY FROM as user with permissions and BYPASSRLS SET SESSION AUTHORIZATION regress_rls_exempt_user; @@ -1649,8 +1651,10 @@ COPY copy_t FROM STDIN; --ok SET SESSION AUTHORIZATION regress_rls_carol; SET row_security TO OFF; COPY copy_t FROM STDIN; --fail - permission denied. +\. SET row_security TO ON; COPY copy_t FROM STDIN; --fail - permission denied. +\. RESET SESSION AUTHORIZATION; DROP TABLE copy_t;