diff options
Diffstat (limited to 'src/bin/psql/common.c')
| -rw-r--r-- | src/bin/psql/common.c | 105 |
1 files changed, 96 insertions, 9 deletions
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); @@ -616,6 +616,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 * * When our command string contained a COPY FROM STDIN or COPY TO STDOUT, @@ -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; } |