diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index 8abadb2..2e8309d 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -748,6 +748,29 @@ exec_command(const char *cmd, status = PSQL_CMD_SEND; } + /* \gset send query and store result */ + else if (strcmp(cmd, "gset") == 0) + { + bool error; + + pset.gvars = psql_scan_slash_vars(scan_state, &error); + + if (!pset.gvars) + { + psql_error("\\%s: missing required argument\n", cmd); + status = PSQL_CMD_NOSEND; + } + else if (error) + { + psql_error("\\%s: syntax error\n", cmd); + status = PSQL_CMD_NOSEND; + tglist_free(pset.gvars); + pset.gvars = NULL; + } + else + status = PSQL_CMD_SEND; + } + /* help */ else if (strcmp(cmd, "h") == 0 || strcmp(cmd, "help") == 0) { diff --git a/src/bin/psql/command.h b/src/bin/psql/command.h index f0bcea0..57cbe23 100644 --- a/src/bin/psql/command.h +++ b/src/bin/psql/command.h @@ -16,6 +16,7 @@ typedef enum _backslashResult { PSQL_CMD_UNKNOWN = 0, /* not done parsing yet (internal only) */ PSQL_CMD_SEND, /* query complete; send off */ + PSQL_CMD_NOSEND, /* query complete, don't send */ PSQL_CMD_SKIP_LINE, /* keep building query */ PSQL_CMD_TERMINATE, /* quit program */ PSQL_CMD_NEWEDIT, /* query buffer was changed (e.g., via \e) */ diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index 330d5ce..5248721 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -826,6 +826,103 @@ PrintQueryResults(PGresult *results) return success; } +/* + * StoreQueryResult: store first row of result to selected variables + * + * Note: Utility function for use by SendQuery() only. + * + * Returns true if the query executed sucessfully, false otherwise. + */ +static bool +StoreQueryResult(PGresult *result) +{ + bool success; + + switch (PQresultStatus(result)) + { + case PGRES_TUPLES_OK: + { + int i; + + if (PQntuples(result) < 1) + { + psql_error("no data found\n"); + success = false; + } + else if (PQntuples(result) > 1) + { + psql_error("too many rows\n"); + success = false; + } + else + { + TargetListData *iter = (TargetListData *) pset.gvars; + + success = true; + + for (i = 0; i < PQnfields(result); i++) + { + if (!iter) + { + psql_error("to few target variables\n"); + success = false; + break; + } + + if (iter->name) + { + if (!SetVariable(pset.vars, iter->name, + PQgetvalue(result, 0, i))) + { + psql_error("invalid variable name: \"%s\"", + iter->name); + success = false; + break; + } + } + + iter = iter->next; + } + + if (success && iter != NULL) + { + psql_error("too many target variables\n"); + success = false; + } + } + } + break; + + case PGRES_COMMAND_OK: + case PGRES_EMPTY_QUERY: + psql_error("no data found\n"); + success = false; + break; + + case PGRES_COPY_OUT: + case PGRES_COPY_IN: + psql_error("COPY isnot supported by \\gset command\n"); + success = false; + break; + + case PGRES_BAD_RESPONSE: + case PGRES_NONFATAL_ERROR: + case PGRES_FATAL_ERROR: + success = false; + break; + + default: + success = false; + psql_error("unexpected PQresultStatus: %d\n", + PQresultStatus(result)); + break; + } + + tglist_free(pset.gvars); + pset.gvars = NULL; + + return success; +} /* * SendQuery: send the query string to the backend @@ -953,7 +1050,12 @@ SendQuery(const char *query) /* but printing results isn't: */ if (OK && results) - OK = PrintQueryResults(results); + { + if (pset.gvars) + OK = StoreQueryResult(results); + else + OK = PrintQueryResults(results); + } } else { @@ -1668,3 +1770,53 @@ expand_tilde(char **filename) return *filename; } + + +/* + * Add name of internal variable to query targer list + * + */ +TargetList +tglist_add(TargetList tglist, const char *name) +{ + TargetListData *tgf; + + tgf = pg_malloc(sizeof(TargetListData)); + tgf->name = name ? pg_strdup(name) : NULL; + tgf->next = NULL; + + if (tglist) + { + TargetListData *iter = (TargetListData *) tglist; + + while (iter->next) + iter = iter->next; + + iter->next = tgf; + + return tglist; + } + else + return (TargetList) tgf; +} + +/* + * Release target list + * + */ +void +tglist_free(TargetList tglist) +{ + TargetListData *iter = (TargetListData *) tglist; + + while (iter) + { + TargetListData *next = iter->next; + + if (iter->name) + free(iter->name); + + free(iter); + iter = next; + } +} diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h index 8037cbc..f955fec 100644 --- a/src/bin/psql/common.h +++ b/src/bin/psql/common.h @@ -21,6 +21,14 @@ #define atooid(x) ((Oid) strtoul((x), NULL, 10)) +typedef struct _target_field +{ + char *name; + struct _target_field *next; +} TargetListData; + +typedef struct TargetListData *TargetList; + /* * Safer versions of some standard C library functions. If an * out-of-memory condition occurs, these functions will bail out @@ -63,4 +71,7 @@ extern const char *session_username(void); extern char *expand_tilde(char **filename); +extern TargetList tglist_add(TargetList tglist, const char *name); +extern void tglist_free(TargetList tglist); + #endif /* COMMON_H */ diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c index 3ebf7cc..0f34f41 100644 --- a/src/bin/psql/help.c +++ b/src/bin/psql/help.c @@ -162,13 +162,14 @@ slashUsage(unsigned short int pager) { FILE *output; - output = PageOutput(94, pager); + output = PageOutput(95, pager); /* if you add/remove a line here, change the row count above */ fprintf(output, _("General\n")); fprintf(output, _(" \\copyright show PostgreSQL usage and distribution terms\n")); fprintf(output, _(" \\g [FILE] or ; execute query (and send results to file or |pipe)\n")); + fprintf(output, _(" \\gset NAME [, NAME [..]] execute query and store result in internal variables\n")); fprintf(output, _(" \\h [NAME] help on syntax of SQL commands, * for all commands\n")); fprintf(output, _(" \\q quit psql\n")); fprintf(output, "\n"); diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c index 16c65ec..33faff8 100644 --- a/src/bin/psql/mainloop.c +++ b/src/bin/psql/mainloop.c @@ -327,6 +327,14 @@ MainLoop(FILE *source) /* flush any paren nesting info after forced send */ psql_scan_reset(scan_state); } + else if (slashCmdStatus == PSQL_CMD_NOSEND) + { + resetPQExpBuffer(query_buf); + /* reset parsing state since we are rescanning whole line */ + psql_scan_reset(scan_state); + line_saved_in_history = true; + success = false; + } else if (slashCmdStatus == PSQL_CMD_NEWEDIT) { /* rescan query_buf as new input */ diff --git a/src/bin/psql/psqlscan.h b/src/bin/psql/psqlscan.h index 0b5ef0c..742de7f 100644 --- a/src/bin/psql/psqlscan.h +++ b/src/bin/psql/psqlscan.h @@ -8,6 +8,7 @@ #ifndef PSQLSCAN_H #define PSQLSCAN_H +#include "common.h" #include "pqexpbuffer.h" #include "prompt.h" @@ -33,7 +34,8 @@ enum slash_option_type OT_SQLIDHACK, /* SQL identifier, but don't downcase */ OT_FILEPIPE, /* it's a filename or pipe */ OT_WHOLE_LINE, /* just snarf the rest of the line */ - OT_NO_EVAL /* no expansion of backticks or variables */ + OT_NO_EVAL, /* no expansion of backticks or variables */ + OT_VARLIST /* returns variable's identifier or comma */ }; @@ -59,6 +61,10 @@ extern char *psql_scan_slash_option(PsqlScanState state, char *quote, bool semicolon); +extern TargetList psql_scan_slash_vars(PsqlScanState state, + bool *error); + + extern void psql_scan_slash_command_end(PsqlScanState state); #endif /* PSQLSCAN_H */ diff --git a/src/bin/psql/psqlscan.l b/src/bin/psql/psqlscan.l index 1208c8f..15018af 100644 --- a/src/bin/psql/psqlscan.l +++ b/src/bin/psql/psqlscan.l @@ -106,13 +106,18 @@ static char *option_quote; static int unquoted_option_chars; static int backtick_start_offset; +TargetList gvars; + /* Return values from yylex() */ #define LEXRES_EOL 0 /* end of input */ #define LEXRES_SEMI 1 /* command-terminating semicolon found */ -#define LEXRES_BACKSLASH 2 /* backslash command start */ -#define LEXRES_OK 3 /* OK completion of backslash argument */ - +#define LEXRES_COMMA 2 +#define LEXRES_BACKSLASH 3 /* backslash command start */ +#define LEXRES_IDENT 4 /* valid internal variable identifier */ +#define LEXRES_OK 5 /* OK completion of backslash argument */ +#define LEXRES_ERROR 6 +#define LEXRES_NONE 7 int yylex(void); @@ -167,6 +172,7 @@ static void escape_variable(bool as_ident); * $foo$ quoted strings * quoted identifier with Unicode escapes * quoted string with Unicode escapes + * comma separated list of variables * * Note: we intentionally don't mimic the backend's state; we have * no need to distinguish it from state, and no good way to get out @@ -183,6 +189,7 @@ static void escape_variable(bool as_ident); %x xdolq %x xui %x xus +%x xvl /* Additional exclusive states for psql only: lex backslash commands */ %x xslashcmd %x xslashargstart @@ -628,6 +635,26 @@ other . ECHO; } +{ + +{identifier} { + gvars = tglist_add(gvars, yytext); + return LEXRES_IDENT; + } + +"," { + return LEXRES_COMMA; + } + +{horiz_space}+ { } + +. | +\n { + return LEXRES_ERROR; + } + +} + {xufailed} { /* throw back all but the initial u/U */ yyless(1); @@ -1449,6 +1476,75 @@ psql_scan_slash_command(PsqlScanState state) } /* + * returns identifier from identifier list + * + * when comma_expected is true, when we require comma before identifier + * error is true, when unexpected char was identified or missing + * comma. + * + */ +TargetList +psql_scan_slash_vars(PsqlScanState state, + bool *error) +{ + PQExpBufferData mybuf; + int lexresult = LEXRES_NONE; + int lexresult_prev; + bool is_lexer_error = false; + + + /* Must be scanning already */ + psql_assert(state->scanbufhandle); + + /* Set up static variables that will be used by yylex */ + cur_state = state; + output_buf = &mybuf; + + if (state->buffer_stack != NULL) + yy_switch_to_buffer(state->buffer_stack->buf); + else + yy_switch_to_buffer(state->scanbufhandle); + + gvars = NULL; + + do + { + lexresult_prev = lexresult; + + BEGIN(xvl); + lexresult = yylex(); + + /* read to end */ + if (is_lexer_error) + continue; + + if (lexresult == LEXRES_ERROR) + { + is_lexer_error = true; + continue; + } + + if (lexresult_prev == LEXRES_IDENT && lexresult == LEXRES_IDENT) + { + is_lexer_error = true; + continue; + } + + if (!is_lexer_error && (lexresult_prev == LEXRES_COMMA || lexresult_prev == LEXRES_NONE) + && lexresult == LEXRES_COMMA) + gvars = tglist_add(gvars, NULL); + + } while (lexresult != LEXRES_EOL); + + if (!is_lexer_error && lexresult_prev == LEXRES_COMMA) + gvars = tglist_add(gvars, NULL); + + *error = is_lexer_error; + + return gvars; +} + +/* * Parse off the next argument for a backslash command, and return it as a * malloc'd string. If there are no more arguments, returns NULL. * diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h index c907fa0..e3381b2 100644 --- a/src/bin/psql/settings.h +++ b/src/bin/psql/settings.h @@ -9,6 +9,7 @@ #define SETTINGS_H +#include "common.h" #include "variables.h" #include "print.h" @@ -73,6 +74,7 @@ typedef struct _psqlSettings printQueryOpt popt; char *gfname; /* one-shot file output argument for \g */ + TargetList gvars; /* one-shot target list argument for \gset */ bool notty; /* stdin or stdout is not a tty (as determined * on startup) */ diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c index a1bb230..b47a44a 100644 --- a/src/bin/psql/tab-complete.c +++ b/src/bin/psql/tab-complete.c @@ -811,7 +811,7 @@ psql_completion(char *text, int start, int end) "\\dF", "\\dFd", "\\dFp", "\\dFt", "\\dg", "\\di", "\\dl", "\\dL", "\\dn", "\\do", "\\dp", "\\drds", "\\ds", "\\dS", "\\dt", "\\dT", "\\dv", "\\du", "\\e", "\\echo", "\\ef", "\\encoding", - "\\f", "\\g", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l", + "\\f", "\\g", "\\gset", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l", "\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink", "\\o", "\\p", "\\password", "\\prompt", "\\pset", "\\q", "\\qecho", "\\r", "\\set", "\\sf", "\\t", "\\T",