diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml index eb5e3b19048..233f747163e 100644 --- a/doc/src/sgml/ref/psql-ref.sgml +++ b/doc/src/sgml/ref/psql-ref.sgml @@ -1620,6 +1620,46 @@ Tue Oct 26 21:40:57 CEST 1999 + + \gset [ prefix ] + + + + Sends the current query input buffer to the server and stores the + query's output into psql variables (see ). + The query to be executed must return exactly one row. Each column of + the row is stored into a separate variable, named the same as the + column. For example: + +=> SELECT 'hello' AS var1, 10 AS var2 +-> \gset +=> \echo :var1 :var2 +hello 10 + + + + If you specify a prefix, + that string is prepended to the query's column names to create the + variable names to use: + +=> SELECT 'hello' AS var1, 10 AS var2 +-> \gset result_ +=> \echo :result_var1 :result_var2 +hello 10 + + + + If a column result is NULL, the corresponding variable is unset + rather than being set. + + + If the query fails or does not return one row, + no variables are changed. + + + + \h or \help [ command ] diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index 1e9aa89089e..012cb75e52c 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -731,7 +731,7 @@ exec_command(const char *cmd, free(fname); } - /* \g [filename] means send query, optionally with output to file/pipe */ + /* \g [filename] -- send query, optionally with output to file/pipe */ else if (strcmp(cmd, "g") == 0) { char *fname = psql_scan_slash_option(scan_state, @@ -748,6 +748,22 @@ exec_command(const char *cmd, status = PSQL_CMD_SEND; } + /* \gset [prefix] -- send query and store result into variables */ + else if (strcmp(cmd, "gset") == 0) + { + char *prefix = psql_scan_slash_option(scan_state, + OT_NORMAL, NULL, false); + + if (prefix) + pset.gset_prefix = prefix; + else + { + /* we must set a non-NULL prefix to trigger storing */ + pset.gset_prefix = pg_strdup(""); + } + status = PSQL_CMD_SEND; + } + /* help */ else if (strcmp(cmd, "h") == 0 || strcmp(cmd, "help") == 0) { diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c index ab517906fca..a8aa1a2df1b 100644 --- a/src/bin/psql/common.c +++ b/src/bin/psql/common.c @@ -43,7 +43,7 @@ pg_strdup(const char *string) if (!string) { psql_error("%s: pg_strdup: cannot duplicate null pointer (internal error)\n", - pset.progname); + pset.progname); exit(EXIT_FAILURE); } tmp = strdup(string); @@ -615,6 +615,65 @@ PrintQueryTuples(const PGresult *results) } +/* + * StoreQueryTuple: assuming query result is OK, save data into variables + * + * Returns true if successful, false otherwise. + */ +static bool +StoreQueryTuple(const PGresult *result) +{ + bool success = true; + + if (PQntuples(result) < 1) + { + psql_error("no rows returned for \\gset\n"); + success = false; + } + else if (PQntuples(result) > 1) + { + psql_error("more than one row returned for \\gset\n"); + success = false; + } + else + { + int i; + + for (i = 0; i < PQnfields(result); i++) + { + char *colname = PQfname(result, i); + char *varname; + char *value; + + /* concate prefix and column name */ + varname = pg_malloc(strlen(pset.gset_prefix) + strlen(colname) + 1); + strcpy(varname, pset.gset_prefix); + strcat(varname, colname); + + if (!PQgetisnull(result, 0, i)) + value = PQgetvalue(result, 0, i); + else + { + /* for NULL value, unset rather than set the variable */ + value = NULL; + } + + if (!SetVariable(pset.vars, varname, value)) + { + psql_error("could not set variable \"%s\"\n", varname); + free(varname); + success = false; + break; + } + + free(varname); + } + } + + return success; +} + + /* * ProcessResult: utility function for use by SendQuery() only * @@ -752,7 +811,7 @@ PrintQueryStatus(PGresult *results) /* - * PrintQueryResults: print out query results as required + * PrintQueryResults: print out (or store) query results as required * * Note: Utility function for use by SendQuery() only. * @@ -770,8 +829,11 @@ PrintQueryResults(PGresult *results) switch (PQresultStatus(results)) { case PGRES_TUPLES_OK: - /* print the data ... */ - success = PrintQueryTuples(results); + /* store or print the data ... */ + if (pset.gset_prefix) + success = StoreQueryTuple(results); + else + success = PrintQueryTuples(results); /* if it's INSERT/UPDATE/DELETE RETURNING, also print status */ cmdstatus = PQcmdStatus(results); if (strncmp(cmdstatus, "INSERT", 6) == 0 || @@ -898,7 +960,7 @@ SendQuery(const char *query) if (on_error_rollback_warning == false && pset.sversion < 80000) { psql_error("The server (version %d.%d) does not support savepoints for ON_ERROR_ROLLBACK.\n", - pset.sversion / 10000, (pset.sversion / 100) % 100); + pset.sversion / 10000, (pset.sversion / 100) % 100); on_error_rollback_warning = true; } else @@ -1046,6 +1108,13 @@ sendquery_cleanup: pset.gfname = NULL; } + /* reset \gset trigger */ + if (pset.gset_prefix) + { + free(pset.gset_prefix); + pset.gset_prefix = NULL; + } + return OK; } @@ -1072,6 +1141,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec) bool started_txn = false; bool did_pager = false; int ntuples; + int fetch_count; char fetch_cmd[64]; instr_time before, after; @@ -1119,9 +1189,18 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec) *elapsed_msec += INSTR_TIME_GET_MILLISEC(after); } + /* + * In \gset mode, we force the fetch count to be 2, so that we will throw + * the appropriate error if the query returns more than one row. + */ + if (pset.gset_prefix) + fetch_count = 2; + else + fetch_count = pset.fetch_count; + snprintf(fetch_cmd, sizeof(fetch_cmd), "FETCH FORWARD %d FROM _psql_cursor", - pset.fetch_count); + fetch_count); /* prepare to write output to \g argument, if any */ if (pset.gfname) @@ -1147,7 +1226,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec) if (pset.timing) INSTR_TIME_SET_CURRENT(before); - /* get FETCH_COUNT tuples at a time */ + /* get fetch_count tuples at a time */ results = PQexec(pset.db, fetch_cmd); if (pset.timing) @@ -1174,9 +1253,17 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec) break; } + if (pset.gset_prefix) + { + /* StoreQueryTuple will complain if not exactly one row */ + OK = StoreQueryTuple(results); + PQclear(results); + break; + } + ntuples = PQntuples(results); - if (ntuples < pset.fetch_count) + if (ntuples < fetch_count) { /* this is the last result set, so allow footer decoration */ my_popt.topt.stop_table = true; @@ -1214,7 +1301,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec) * writing things to the stream, we presume $PAGER has disappeared and * stop bothering to pull down more data. */ - if (ntuples < pset.fetch_count || cancel_pressed || flush_error || + if (ntuples < fetch_count || cancel_pressed || flush_error || ferror(pset.queryFout)) break; } diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c index fd7effa87b3..43cb550bd24 100644 --- a/src/bin/psql/help.c +++ b/src/bin/psql/help.c @@ -165,13 +165,14 @@ slashUsage(unsigned short int pager) currdb = PQdb(pset.db); - 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 [PREFIX] execute query and store results in psql 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"); @@ -261,7 +262,7 @@ slashUsage(unsigned short int pager) currdb); else fprintf(output, _(" \\c[onnect] [DBNAME|- USER|- HOST|- PORT|-]\n" - " connect to new database (currently no connection)\n")), + " connect to new database (currently no connection)\n")); fprintf(output, _(" \\encoding [ENCODING] show or set client encoding\n")); fprintf(output, _(" \\password [USERNAME] securely change the password for a user\n")); fprintf(output, _(" \\conninfo display information about current connection\n")); diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h index a80fb990c16..e78aa9a4999 100644 --- a/src/bin/psql/settings.h +++ b/src/bin/psql/settings.h @@ -73,6 +73,7 @@ typedef struct _psqlSettings printQueryOpt popt; char *gfname; /* one-shot file output argument for \g */ + char *gset_prefix; /* one-shot prefix 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 f7d84669f97..09396ca5900 100644 --- a/src/bin/psql/tab-complete.c +++ b/src/bin/psql/tab-complete.c @@ -856,7 +856,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", diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out new file mode 100644 index 00000000000..7dd5c7b51b9 --- /dev/null +++ b/src/test/regress/expected/psql.out @@ -0,0 +1,54 @@ +-- +-- Tests for psql features that aren't closely connected to any +-- specific server features +-- +-- \gset +select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_ +\echo :pref01_test01 :pref01_test02 :pref01_test03 +10 20 Hello +-- should fail: bad variable name +select 10 as "bad name" +\gset +could not set variable "bad name" +-- multiple backslash commands in one line +select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x +1 +select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y +3 +4 +select 5 as x, 6 as y \gset pref01_ \\ \g \echo :pref01_x :pref01_y + x | y +---+--- + 5 | 6 +(1 row) + +5 6 +select 7 as x, 8 as y \g \gset pref01_ \echo :pref01_x :pref01_y + x | y +---+--- + 7 | 8 +(1 row) + +7 8 +-- NULL should unset the variable +\set var2 xyz +select 1 as var1, NULL as var2, 3 as var3 \gset +\echo :var1 :var2 :var3 +1 :var2 3 +-- \gset requires just one tuple +select 10 as test01, 20 as test02 from generate_series(1,3) \gset +more than one row returned for \gset +select 10 as test01, 20 as test02 from generate_series(1,0) \gset +no rows returned for \gset +-- \gset should work in FETCH_COUNT mode too +\set FETCH_COUNT 1 +select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x +1 +select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y +3 +4 +select 10 as test01, 20 as test02 from generate_series(1,3) \gset +more than one row returned for \gset +select 10 as test01, 20 as test02 from generate_series(1,0) \gset +no rows returned for \gset +\unset FETCH_COUNT diff --git a/src/test/regress/parallel_schedule b/src/test/regress/parallel_schedule index bdcf3a6a559..d3def07f92c 100644 --- a/src/test/regress/parallel_schedule +++ b/src/test/regress/parallel_schedule @@ -88,7 +88,7 @@ test: privileges security_label collate # ---------- # Another group of parallel tests # ---------- -test: misc alter_generic +test: alter_generic misc psql # rules cannot run concurrently with any test that creates a view test: rules diff --git a/src/test/regress/serial_schedule b/src/test/regress/serial_schedule index c7c2ed0f6a0..7059fca092b 100644 --- a/src/test/regress/serial_schedule +++ b/src/test/regress/serial_schedule @@ -95,8 +95,9 @@ test: prepared_xacts test: privileges test: security_label test: collate -test: misc test: alter_generic +test: misc +test: psql test: rules test: event_trigger test: select_views diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql new file mode 100644 index 00000000000..50ee02a7f79 --- /dev/null +++ b/src/test/regress/sql/psql.sql @@ -0,0 +1,39 @@ +-- +-- Tests for psql features that aren't closely connected to any +-- specific server features +-- + +-- \gset + +select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_ + +\echo :pref01_test01 :pref01_test02 :pref01_test03 + +-- should fail: bad variable name +select 10 as "bad name" +\gset + +-- multiple backslash commands in one line +select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x +select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y +select 5 as x, 6 as y \gset pref01_ \\ \g \echo :pref01_x :pref01_y +select 7 as x, 8 as y \g \gset pref01_ \echo :pref01_x :pref01_y + +-- NULL should unset the variable +\set var2 xyz +select 1 as var1, NULL as var2, 3 as var3 \gset +\echo :var1 :var2 :var3 + +-- \gset requires just one tuple +select 10 as test01, 20 as test02 from generate_series(1,3) \gset +select 10 as test01, 20 as test02 from generate_series(1,0) \gset + +-- \gset should work in FETCH_COUNT mode too +\set FETCH_COUNT 1 + +select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x +select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y +select 10 as test01, 20 as test02 from generate_series(1,3) \gset +select 10 as test01, 20 as test02 from generate_series(1,0) \gset + +\unset FETCH_COUNT