summaryrefslogtreecommitdiffstats
path: root/src/fe_utils/connect_utils.c
blob: 7a1edea7c8cd9f3d49987e2c46ad17c34158f757 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
/*-------------------------------------------------------------------------
 *
 * Facilities for frontend code to connect to and disconnect from databases.
 *
 * Portions Copyright (c) 1996-2023, 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_fatal("could not connect to database %s: out of memory",
					 cparams->dbname);

		/*
		 * No luck?  Trying asking (again) for a password.
		 */
		if (PQstatus(conn) == CONNECTION_BAD &&
			PQconnectionNeedsPassword(conn) &&
			cparams->prompt_password != TRI_NO)
		{
			PQfinish(conn);
			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_fatal("%s", PQerrorMessage(conn));
	}

	/* 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);
}