diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c index 596d112..8793fd2 100644 --- a/src/bin/pgbench/pgbench.c +++ b/src/bin/pgbench/pgbench.c @@ -54,6 +54,7 @@ #endif #include "pgbench.h" +#include "psqlscan.h" #define ERRCODE_UNDEFINED_TABLE "42P01" @@ -285,7 +286,7 @@ typedef enum QueryMode static QueryMode querymode = QUERY_SIMPLE; static const char *QUERYMODE[] = {"simple", "extended", "prepared"}; -typedef struct +typedef struct Command_t { char *line; /* full text of command line */ int command_num; /* unique index of this Command struct */ @@ -295,6 +296,7 @@ typedef struct int cols[MAX_ARGS]; /* corresponding column starting from 1 */ PgBenchExpr *expr; /* parsed expression */ SimpleStats stats; /* time spent in this command */ + struct Command_t *next; /* more command if any, for multistatements */ } Command; static struct @@ -303,6 +305,22 @@ static struct Command **commands; StatsData stats; } sql_script[MAX_SCRIPTS]; /* SQL script files */ + +typedef enum +{ + PS_IDLE, + PS_IN_STATEMENT, + PS_IN_BACKSLASH_CMD +} ParseState; + +typedef struct ParseInfo +{ + PsqlScanState scan_state; + PQExpBuffer outbuf; + ParseState mode; +} ParseInfoData; +typedef ParseInfoData *ParseInfo; + static int num_scripts; /* number of scripts in sql_script[] */ static int num_commands = 0; /* total number of Command structs */ static int debug = 0; /* debug flag */ @@ -430,6 +448,9 @@ usage(void) progname, progname); } +PsqlScanCallbacks pgbench_scan_callbacks = +{NULL, NULL, NULL}; + /* * strtoint64 -- convert a string to 64-bit integer * @@ -2287,43 +2308,109 @@ syntax_error(const char *source, const int lineno, exit(1); } -/* Parse a command; return a Command struct, or NULL if it's a comment */ +static ParseInfo +createParseInfo(void) +{ + ParseInfo ret = (ParseInfo) pg_malloc(sizeof(ParseInfoData)); + + ret->scan_state = psql_scan_create(); + ret->outbuf = createPQExpBuffer(); + ret->mode = PS_IDLE; + + return ret; +} + +#define parse_reset_outbuf(pcs) resetPQExpBuffer((pcs)->outbuf) +#define parse_finish_scan(pcs) psql_scan_finish((pcs)->scan_state) + +/* copy a string after removing newlines and collapsing whitespaces */ +static char * +strdup_nonl(const char *in) +{ + char *ret, *p, *q; + + ret = pg_strdup(in); + + /* Replace newlines into spaces */ + for (p = ret ; *p ; p++) + if (*p == '\n') *p = ' '; + + /* collapse successive spaces */ + for (p = q = ret ; *p ; p++, q++) + { + while (isspace(*p) && isspace(*(p + 1))) p++; + if (p > q) *q = *p; + } + *q = '\0'; + + return ret; +} + +/* Parse a backslash command; return a Command struct */ static Command * -process_commands(char *buf, const char *source, const int lineno) +process_backslash_commands(ParseInfo proc_state, char *buf, + const char *source, const int lineno) { const char delim[] = " \f\n\r\t\v"; Command *my_commands; int j; char *p, + *start, *tok; - - /* Make the string buf end at the next newline */ - if ((p = strchr(buf, '\n')) != NULL) - *p = '\0'; + int max_args = -1; /* Skip leading whitespace */ p = buf; while (isspace((unsigned char) *p)) p++; + start = p; - /* If the line is empty or actually a comment, we're done */ - if (*p == '\0' || strncmp(p, "--", 2) == 0) + if (proc_state->mode != PS_IN_BACKSLASH_CMD) + { + if (*p != '\\') + return NULL; /* not a backslash command */ + + /* This is the first line of a backslash command */ + proc_state->mode = PS_IN_BACKSLASH_CMD; + } + + /* + * Make the string buf end at the next newline, or move to just after the + * end of line + */ + if ((p = strchr(start, '\n')) != NULL) + *p = '\0'; + else + p = start + strlen(start); + + /* continued line ends with a backslash */ + if (*(--p) == '\\') + { + *p-- = '\0'; + appendPQExpBufferStr(proc_state->outbuf, start); + + /* Add a delimiter at the end of the line if necessary */ + if (!isspace(*p)) + appendPQExpBufferChar(proc_state->outbuf, ' '); return NULL; + } + + appendPQExpBufferStr(proc_state->outbuf, start); + proc_state->mode = PS_IDLE; + + /* Start parsing the backslash command */ + + p = proc_state->outbuf->data; /* Allocate and initialize Command structure */ my_commands = (Command *) pg_malloc(sizeof(Command)); - my_commands->line = pg_strdup(buf); + my_commands->line = pg_strdup(p); my_commands->command_num = num_commands++; - my_commands->type = 0; /* until set */ + my_commands->type = META_COMMAND; my_commands->argc = 0; + my_commands->next = NULL; initSimpleStats(&my_commands->stats); - if (*p == '\\') - { - int max_args = -1; - - my_commands->type = META_COMMAND; - j = 0; tok = strtok(++p, delim); @@ -2340,6 +2427,7 @@ process_commands(char *buf, const char *source, const int lineno) else tok = strtok(NULL, delim); } + parse_reset_outbuf(proc_state); if (pg_strcasecmp(my_commands->argv[0], "setrandom") == 0) { @@ -2475,28 +2563,91 @@ process_commands(char *buf, const char *source, const int lineno) syntax_error(source, lineno, my_commands->line, my_commands->argv[0], "invalid command", NULL, -1); } + + return my_commands; +} + +/* Parse an input line, return non-null if any command terminates. */ +static Command * +process_commands(ParseInfo proc_state, char *buf, + const char *source, const int lineno) +{ + Command *command = NULL; + Command *retcomd = NULL; + PsqlScanState scan_state = proc_state->scan_state; + promptStatus_t prompt_status = PROMPT_READY; /* dummy */ + PQExpBuffer qbuf = proc_state->outbuf; + PsqlScanResult scan_result; + + if (proc_state->mode != PS_IN_STATEMENT) + { + command = process_backslash_commands(proc_state, buf, source, lineno); + + /* go to next line for continuation of the backslash command. */ + if (command != NULL || proc_state->mode == PS_IN_BACKSLASH_CMD) + return command; } + + /* Parse statements */ + psql_scan_setup(scan_state, buf, strlen(buf), &pgbench_scan_callbacks); + +next_command: + scan_result = psql_scan(scan_state, qbuf, &prompt_status); + + if (scan_result == PSCAN_SEMICOLON) + { + proc_state->mode = PS_IDLE; + /* + * Command is terminated. Fill the struct. + */ + command = (Command*) pg_malloc(sizeof(Command)); + command->line = strdup_nonl(qbuf->data); + command->command_num = num_commands++; + command->type = SQL_COMMAND; + command->argc = 0; + command->next = NULL; + + /* Put this command at the end of returning command chain */ + if (!retcomd) + retcomd = command; else { - my_commands->type = SQL_COMMAND; + Command *pcomm = retcomd; + while (pcomm->next) pcomm = pcomm->next; + pcomm->next = command; + } switch (querymode) { case QUERY_SIMPLE: - my_commands->argv[0] = pg_strdup(p); - my_commands->argc++; + command->argv[0] = pg_strdup(qbuf->data); + command->argc++; break; case QUERY_EXTENDED: case QUERY_PREPARED: - if (!parseQuery(my_commands, p)) + if (!parseQuery(command, qbuf->data)) exit(1); break; default: exit(1); } + + parse_reset_outbuf(proc_state); + + /* Ask for the next statement in this line */ + goto next_command; } + else if (scan_result == PSCAN_BACKSLASH) + { + fprintf(stderr, "Unexpected backslash in SQL statement: %s:%d\n", + source, lineno); + exit(1); + } + + proc_state->mode = PS_IN_STATEMENT; + psql_scan_finish(scan_state); - return my_commands; + return retcomd; } /* @@ -2557,6 +2708,7 @@ process_file(char *filename) index; char *buf; int alloc_num; + ParseInfo proc_state = createParseInfo(); alloc_num = COMMANDS_ALLOC_NUM; my_commands = (Command **) pg_malloc(sizeof(Command *) * alloc_num); @@ -2571,26 +2723,38 @@ process_file(char *filename) return NULL; } + proc_state->mode = PS_IDLE; + lineno = 0; index = 0; while ((buf = read_line_from_file(fd)) != NULL) { - Command *command; + Command *command = NULL; lineno += 1; - command = process_commands(buf, filename, lineno); + command = process_commands(proc_state, buf, filename, lineno); free(buf); if (command == NULL) + { + /* + * command is NULL when psql_scan returns PSCAN_EOL or + * PSCAN_INCOMPLETE. Immediately ask for the next line for the + * cases. + */ continue; + } - my_commands[index] = command; - index++; + while (command) + { + my_commands[index++] = command; + command = command->next; + } - if (index >= alloc_num) + if (index > alloc_num) { alloc_num += COMMANDS_ALLOC_NUM; my_commands = pg_realloc(my_commands, sizeof(Command *) * alloc_num); @@ -2598,6 +2762,8 @@ process_file(char *filename) } fclose(fd); + parse_finish_scan(proc_state); + my_commands[index] = NULL; return my_commands; @@ -2613,6 +2779,7 @@ process_builtin(const char *tb, const char *source) index; char buf[BUFSIZ]; int alloc_num; + ParseInfo proc_state = createParseInfo(); alloc_num = COMMANDS_ALLOC_NUM; my_commands = (Command **) pg_malloc(sizeof(Command *) * alloc_num); @@ -2639,10 +2806,12 @@ process_builtin(const char *tb, const char *source) lineno += 1; - command = process_commands(buf, source, lineno); + command = process_commands(proc_state, buf, source, lineno); if (command == NULL) continue; + /* builtin doesn't need multistatements */ + Assert(command->next == NULL); my_commands[index] = command; index++; @@ -2654,6 +2823,7 @@ process_builtin(const char *tb, const char *source) } my_commands[index] = NULL; + parse_finish_scan(proc_state); return my_commands; }