Create a psql command \gset to store query results into psql variables.

This eases manipulation of query results in psql scripts.

Pavel Stehule, reviewed by Piyush Newe, Shigeru Hanada, and Tom Lane
This commit is contained in:
Tom Lane 2013-02-02 17:06:38 -05:00
parent 101d6ae755
commit d2d153fdb0
10 changed files with 254 additions and 15 deletions

View File

@ -1620,6 +1620,46 @@ Tue Oct 26 21:40:57 CEST 1999
</listitem>
</varlistentry>
<varlistentry>
<term><literal>\gset</literal> [ <replaceable class="parameter">prefix</replaceable> ]</term>
<listitem>
<para>
Sends the current query input buffer to the server and stores the
query's output into <application>psql</> variables (see <xref
linkend="APP-PSQL-variables" endterm="APP-PSQL-variables-title">).
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:
<programlisting>
=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
-&gt; <userinput>\gset</userinput>
=&gt; <userinput>\echo :var1 :var2</userinput>
hello 10
</programlisting>
</para>
<para>
If you specify a <replaceable class="parameter">prefix</replaceable>,
that string is prepended to the query's column names to create the
variable names to use:
<programlisting>
=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
-&gt; <userinput>\gset result_</userinput>
=&gt; <userinput>\echo :result_var1 :result_var2</userinput>
hello 10
</programlisting>
</para>
<para>
If a column result is NULL, the corresponding variable is unset
rather than being set.
</para>
<para>
If the query fails or does not return one row,
no variables are changed.
</para>
</listitem>
</varlistentry>
<varlistentry>
<term><literal>\h</literal> or <literal>\help</literal> <literal>[ <replaceable class="parameter">command</replaceable> ]</literal></term>
<listitem>

View File

@ -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)
{

View File

@ -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;
}

View File

@ -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"));

View File

@ -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) */

View File

@ -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",

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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