diff options
author | Robert Haas <rhaas@postgresql.org> | 2021-02-05 13:33:38 -0500 |
---|---|---|
committer | Robert Haas <rhaas@postgresql.org> | 2021-02-05 13:33:38 -0500 |
commit | e955bd4b6c2bcdbd253837f6cf4c7520b98e69d4 (patch) | |
tree | bc8e70ed9fbe9bee458299dc6f76b20311185cf7 /src/fe_utils/connect_utils.c | |
parent | c444472af5c202067a9ecb0ff8df7370fb1ea8f4 (diff) | |
download | postgresql-e955bd4b6c2bcdbd253837f6cf4c7520b98e69d4.tar.gz postgresql-e955bd4b6c2bcdbd253837f6cf4c7520b98e69d4.zip |
Move some code from src/bin/scripts to src/fe_utils to permit reuse.
The parallel slots infrastructure (which implements client-side
multiplexing of server connections doing similar things, not
threading or multiple processes or anything like that) are moved from
src/bin/scripts/scripts_parallel.c to src/fe_utils/parallel_slot.c.
The functions consumeQueryResult() and processQueryResult() which were
previously part of src/bin/scripts/common.c are now moved into that
file as well, becoming static helper functions. This might need to be
changed in the future, but currently they're not used for anything
else.
Some other functions from src/bin/scripts/common.c are moved to to
src/fe_utils and are split up among several files. connectDatabase(),
connectMaintenanceDatabase(), and disconnectDatabase() are moved to
connect_utils.c. executeQuery(), executeCommand(), and
executeMaintenanceCommand() are move to query_utils.c.
handle_help_version_opts() is moved to option_utils.c.
Mark Dilger, reviewed by me. The larger patch series of which this is
a part has also had review from Peter Geoghegan, Andres Freund, Álvaro
Herrera, Michael Paquier, and Amul Sul, but I don't know whether any
of them have reviewed this bit specifically.
Discussion: http://postgr.es/m/12ED3DA8-25F0-4B68-937D-D907CFBF08E7@enterprisedb.com
Discussion: http://postgr.es/m/5F743835-3399-419C-8324-2D424237E999@enterprisedb.com
Discussion: http://postgr.es/m/70655DF3-33CE-4527-9A4D-DDEB582B6BA0@enterprisedb.com
Diffstat (limited to 'src/fe_utils/connect_utils.c')
-rw-r--r-- | src/fe_utils/connect_utils.c | 181 |
1 files changed, 181 insertions, 0 deletions
diff --git a/src/fe_utils/connect_utils.c b/src/fe_utils/connect_utils.c new file mode 100644 index 00000000000..96bb798316a --- /dev/null +++ b/src/fe_utils/connect_utils.c @@ -0,0 +1,181 @@ +/*------------------------------------------------------------------------- + * + * Facilities for frontend code to connect to and disconnect from databases. + * + * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * src/fe_utils/connect_utils.c + * + *------------------------------------------------------------------------- + */ +#include "postgres_fe.h" + +#include "common/connect.h" +#include "common/logging.h" +#include "common/string.h" +#include "fe_utils/connect_utils.h" +#include "fe_utils/query_utils.h" + +/* + * Make a database connection with the given parameters. + * + * An interactive password prompt is automatically issued if needed and + * allowed by cparams->prompt_password. + * + * If allow_password_reuse is true, we will try to re-use any password + * given during previous calls to this routine. (Callers should not pass + * allow_password_reuse=true unless reconnecting to the same database+user + * as before, else we might create password exposure hazards.) + */ +PGconn * +connectDatabase(const ConnParams *cparams, const char *progname, + bool echo, bool fail_ok, bool allow_password_reuse) +{ + PGconn *conn; + bool new_pass; + static char *password = NULL; + + /* Callers must supply at least dbname; other params can be NULL */ + Assert(cparams->dbname); + + if (!allow_password_reuse && password) + { + free(password); + password = NULL; + } + + if (cparams->prompt_password == TRI_YES && password == NULL) + password = simple_prompt("Password: ", false); + + /* + * Start the connection. Loop until we have a password if requested by + * backend. + */ + do + { + const char *keywords[8]; + const char *values[8]; + int i = 0; + + /* + * If dbname is a connstring, its entries can override the other + * values obtained from cparams; but in turn, override_dbname can + * override the dbname component of it. + */ + keywords[i] = "host"; + values[i++] = cparams->pghost; + keywords[i] = "port"; + values[i++] = cparams->pgport; + keywords[i] = "user"; + values[i++] = cparams->pguser; + keywords[i] = "password"; + values[i++] = password; + keywords[i] = "dbname"; + values[i++] = cparams->dbname; + if (cparams->override_dbname) + { + keywords[i] = "dbname"; + values[i++] = cparams->override_dbname; + } + keywords[i] = "fallback_application_name"; + values[i++] = progname; + keywords[i] = NULL; + values[i++] = NULL; + Assert(i <= lengthof(keywords)); + + new_pass = false; + conn = PQconnectdbParams(keywords, values, true); + + if (!conn) + { + pg_log_error("could not connect to database %s: out of memory", + cparams->dbname); + exit(1); + } + + /* + * No luck? Trying asking (again) for a password. + */ + if (PQstatus(conn) == CONNECTION_BAD && + PQconnectionNeedsPassword(conn) && + cparams->prompt_password != TRI_NO) + { + PQfinish(conn); + if (password) + free(password); + password = simple_prompt("Password: ", false); + new_pass = true; + } + } while (new_pass); + + /* check to see that the backend connection was successfully made */ + if (PQstatus(conn) == CONNECTION_BAD) + { + if (fail_ok) + { + PQfinish(conn); + return NULL; + } + pg_log_error("%s", PQerrorMessage(conn)); + exit(1); + } + + /* Start strict; callers may override this. */ + PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL, echo)); + + return conn; +} + +/* + * Try to connect to the appropriate maintenance database. + * + * This differs from connectDatabase only in that it has a rule for + * inserting a default "dbname" if none was given (which is why cparams + * is not const). Note that cparams->dbname should typically come from + * a --maintenance-db command line parameter. + */ +PGconn * +connectMaintenanceDatabase(ConnParams *cparams, + const char *progname, bool echo) +{ + PGconn *conn; + + /* If a maintenance database name was specified, just connect to it. */ + if (cparams->dbname) + return connectDatabase(cparams, progname, echo, false, false); + + /* Otherwise, try postgres first and then template1. */ + cparams->dbname = "postgres"; + conn = connectDatabase(cparams, progname, echo, true, false); + if (!conn) + { + cparams->dbname = "template1"; + conn = connectDatabase(cparams, progname, echo, false, false); + } + return conn; +} + +/* + * Disconnect the given connection, canceling any statement if one is active. + */ +void +disconnectDatabase(PGconn *conn) +{ + char errbuf[256]; + + Assert(conn != NULL); + + if (PQtransactionStatus(conn) == PQTRANS_ACTIVE) + { + PGcancel *cancel; + + if ((cancel = PQgetCancel(conn))) + { + (void) PQcancel(cancel, errbuf, sizeof(errbuf)); + PQfreeCancel(cancel); + } + } + + PQfinish(conn); +} |