1/*-------------------------------------------------------------------------
2 *
3 * common.c
4 * Common support routines for bin/scripts/
5 *
6 *
7 * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
8 * Portions Copyright (c) 1994, Regents of the University of California
9 *
10 * src/bin/scripts/common.c
11 *
12 *-------------------------------------------------------------------------
13 */
14
15#include "postgres_fe.h"
16
17#include <signal.h>
18#include <unistd.h>
19
20#include "common.h"
21#include "common/logging.h"
22#include "fe_utils/connect.h"
23#include "fe_utils/string_utils.h"
24
25
26static PGcancel *volatile cancelConn = NULL;
27bool CancelRequested = false;
28
29#ifdef WIN32
30static CRITICAL_SECTION cancelConnLock;
31#endif
32
33/*
34 * Provide strictly harmonized handling of --help and --version
35 * options.
36 */
37void
38handle_help_version_opts(int argc, char *argv[],
39 const char *fixed_progname, help_handler hlp)
40{
41 if (argc > 1)
42 {
43 if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
44 {
45 hlp(get_progname(argv[0]));
46 exit(0);
47 }
48 if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
49 {
50 printf("%s (PostgreSQL) " PG_VERSION "\n", fixed_progname);
51 exit(0);
52 }
53 }
54}
55
56
57/*
58 * Make a database connection with the given parameters.
59 *
60 * An interactive password prompt is automatically issued if needed and
61 * allowed by prompt_password.
62 *
63 * If allow_password_reuse is true, we will try to re-use any password
64 * given during previous calls to this routine. (Callers should not pass
65 * allow_password_reuse=true unless reconnecting to the same database+user
66 * as before, else we might create password exposure hazards.)
67 */
68PGconn *
69connectDatabase(const char *dbname, const char *pghost,
70 const char *pgport, const char *pguser,
71 enum trivalue prompt_password, const char *progname,
72 bool echo, bool fail_ok, bool allow_password_reuse)
73{
74 PGconn *conn;
75 bool new_pass;
76 static bool have_password = false;
77 static char password[100];
78
79 if (!allow_password_reuse)
80 have_password = false;
81
82 if (!have_password && prompt_password == TRI_YES)
83 {
84 simple_prompt("Password: ", password, sizeof(password), false);
85 have_password = true;
86 }
87
88 /*
89 * Start the connection. Loop until we have a password if requested by
90 * backend.
91 */
92 do
93 {
94 const char *keywords[7];
95 const char *values[7];
96
97 keywords[0] = "host";
98 values[0] = pghost;
99 keywords[1] = "port";
100 values[1] = pgport;
101 keywords[2] = "user";
102 values[2] = pguser;
103 keywords[3] = "password";
104 values[3] = have_password ? password : NULL;
105 keywords[4] = "dbname";
106 values[4] = dbname;
107 keywords[5] = "fallback_application_name";
108 values[5] = progname;
109 keywords[6] = NULL;
110 values[6] = NULL;
111
112 new_pass = false;
113 conn = PQconnectdbParams(keywords, values, true);
114
115 if (!conn)
116 {
117 pg_log_error("could not connect to database %s: out of memory",
118 dbname);
119 exit(1);
120 }
121
122 /*
123 * No luck? Trying asking (again) for a password.
124 */
125 if (PQstatus(conn) == CONNECTION_BAD &&
126 PQconnectionNeedsPassword(conn) &&
127 prompt_password != TRI_NO)
128 {
129 PQfinish(conn);
130 simple_prompt("Password: ", password, sizeof(password), false);
131 have_password = true;
132 new_pass = true;
133 }
134 } while (new_pass);
135
136 /* check to see that the backend connection was successfully made */
137 if (PQstatus(conn) == CONNECTION_BAD)
138 {
139 if (fail_ok)
140 {
141 PQfinish(conn);
142 return NULL;
143 }
144 pg_log_error("could not connect to database %s: %s",
145 dbname, PQerrorMessage(conn));
146 exit(1);
147 }
148
149 PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL,
150 progname, echo));
151
152 return conn;
153}
154
155/*
156 * Try to connect to the appropriate maintenance database.
157 */
158PGconn *
159connectMaintenanceDatabase(const char *maintenance_db,
160 const char *pghost, const char *pgport,
161 const char *pguser, enum trivalue prompt_password,
162 const char *progname, bool echo)
163{
164 PGconn *conn;
165
166 /* If a maintenance database name was specified, just connect to it. */
167 if (maintenance_db)
168 return connectDatabase(maintenance_db, pghost, pgport, pguser,
169 prompt_password, progname, echo, false, false);
170
171 /* Otherwise, try postgres first and then template1. */
172 conn = connectDatabase("postgres", pghost, pgport, pguser, prompt_password,
173 progname, echo, true, false);
174 if (!conn)
175 conn = connectDatabase("template1", pghost, pgport, pguser,
176 prompt_password, progname, echo, false, false);
177
178 return conn;
179}
180
181/*
182 * Run a query, return the results, exit program on failure.
183 */
184PGresult *
185executeQuery(PGconn *conn, const char *query, const char *progname, bool echo)
186{
187 PGresult *res;
188
189 if (echo)
190 printf("%s\n", query);
191
192 res = PQexec(conn, query);
193 if (!res ||
194 PQresultStatus(res) != PGRES_TUPLES_OK)
195 {
196 pg_log_error("query failed: %s", PQerrorMessage(conn));
197 pg_log_info("query was: %s", query);
198 PQfinish(conn);
199 exit(1);
200 }
201
202 return res;
203}
204
205
206/*
207 * As above for a SQL command (which returns nothing).
208 */
209void
210executeCommand(PGconn *conn, const char *query,
211 const char *progname, bool echo)
212{
213 PGresult *res;
214
215 if (echo)
216 printf("%s\n", query);
217
218 res = PQexec(conn, query);
219 if (!res ||
220 PQresultStatus(res) != PGRES_COMMAND_OK)
221 {
222 pg_log_error("query failed: %s", PQerrorMessage(conn));
223 pg_log_info("query was: %s", query);
224 PQfinish(conn);
225 exit(1);
226 }
227
228 PQclear(res);
229}
230
231
232/*
233 * As above for a SQL maintenance command (returns command success).
234 * Command is executed with a cancel handler set, so Ctrl-C can
235 * interrupt it.
236 */
237bool
238executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
239{
240 PGresult *res;
241 bool r;
242
243 if (echo)
244 printf("%s\n", query);
245
246 SetCancelConn(conn);
247 res = PQexec(conn, query);
248 ResetCancelConn();
249
250 r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);
251
252 if (res)
253 PQclear(res);
254
255 return r;
256}
257
258
259/*
260 * Split TABLE[(COLUMNS)] into TABLE and [(COLUMNS)] portions. When you
261 * finish using them, pg_free(*table). *columns is a pointer into "spec",
262 * possibly to its NUL terminator.
263 */
264void
265splitTableColumnsSpec(const char *spec, int encoding,
266 char **table, const char **columns)
267{
268 bool inquotes = false;
269 const char *cp = spec;
270
271 /*
272 * Find the first '(' not identifier-quoted. Based on
273 * dequote_downcase_identifier().
274 */
275 while (*cp && (*cp != '(' || inquotes))
276 {
277 if (*cp == '"')
278 {
279 if (inquotes && cp[1] == '"')
280 cp++; /* pair does not affect quoting */
281 else
282 inquotes = !inquotes;
283 cp++;
284 }
285 else
286 cp += PQmblen(cp, encoding);
287 }
288 *table = pg_strdup(spec);
289 (*table)[cp - spec] = '\0'; /* no strndup */
290 *columns = cp;
291}
292
293/*
294 * Break apart TABLE[(COLUMNS)] of "spec". With the reset_val of search_path
295 * in effect, have regclassin() interpret the TABLE portion. Append to "buf"
296 * the qualified name of TABLE, followed by any (COLUMNS). Exit on failure.
297 * We use this to interpret --table=foo under the search path psql would get,
298 * in advance of "ANALYZE public.foo" under the always-secure search path.
299 */
300void
301appendQualifiedRelation(PQExpBuffer buf, const char *spec,
302 PGconn *conn, const char *progname, bool echo)
303{
304 char *table;
305 const char *columns;
306 PQExpBufferData sql;
307 PGresult *res;
308 int ntups;
309
310 splitTableColumnsSpec(spec, PQclientEncoding(conn), &table, &columns);
311
312 /*
313 * Query must remain ABSOLUTELY devoid of unqualified names. This would
314 * be unnecessary given a regclassin() variant taking a search_path
315 * argument.
316 */
317 initPQExpBuffer(&sql);
318 appendPQExpBufferStr(&sql,
319 "SELECT c.relname, ns.nspname\n"
320 " FROM pg_catalog.pg_class c,"
321 " pg_catalog.pg_namespace ns\n"
322 " WHERE c.relnamespace OPERATOR(pg_catalog.=) ns.oid\n"
323 " AND c.oid OPERATOR(pg_catalog.=) ");
324 appendStringLiteralConn(&sql, table, conn);
325 appendPQExpBufferStr(&sql, "::pg_catalog.regclass;");
326
327 executeCommand(conn, "RESET search_path;", progname, echo);
328
329 /*
330 * One row is a typical result, as is a nonexistent relation ERROR.
331 * regclassin() unconditionally accepts all-digits input as an OID; if no
332 * relation has that OID; this query returns no rows. Catalog corruption
333 * might elicit other row counts.
334 */
335 res = executeQuery(conn, sql.data, progname, echo);
336 ntups = PQntuples(res);
337 if (ntups != 1)
338 {
339 pg_log_error(ngettext("query returned %d row instead of one: %s",
340 "query returned %d rows instead of one: %s",
341 ntups),
342 ntups, sql.data);
343 PQfinish(conn);
344 exit(1);
345 }
346 appendPQExpBufferStr(buf,
347 fmtQualifiedId(PQgetvalue(res, 0, 1),
348 PQgetvalue(res, 0, 0)));
349 appendPQExpBufferStr(buf, columns);
350 PQclear(res);
351 termPQExpBuffer(&sql);
352 pg_free(table);
353
354 PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL,
355 progname, echo));
356}
357
358
359/*
360 * Check yes/no answer in a localized way. 1=yes, 0=no, -1=neither.
361 */
362
363/* translator: abbreviation for "yes" */
364#define PG_YESLETTER gettext_noop("y")
365/* translator: abbreviation for "no" */
366#define PG_NOLETTER gettext_noop("n")
367
368bool
369yesno_prompt(const char *question)
370{
371 char prompt[256];
372
373 /*------
374 translator: This is a question followed by the translated options for
375 "yes" and "no". */
376 snprintf(prompt, sizeof(prompt), _("%s (%s/%s) "),
377 _(question), _(PG_YESLETTER), _(PG_NOLETTER));
378
379 for (;;)
380 {
381 char resp[10];
382
383 simple_prompt(prompt, resp, sizeof(resp), true);
384
385 if (strcmp(resp, _(PG_YESLETTER)) == 0)
386 return true;
387 if (strcmp(resp, _(PG_NOLETTER)) == 0)
388 return false;
389
390 printf(_("Please answer \"%s\" or \"%s\".\n"),
391 _(PG_YESLETTER), _(PG_NOLETTER));
392 }
393}
394
395/*
396 * SetCancelConn
397 *
398 * Set cancelConn to point to the current database connection.
399 */
400void
401SetCancelConn(PGconn *conn)
402{
403 PGcancel *oldCancelConn;
404
405#ifdef WIN32
406 EnterCriticalSection(&cancelConnLock);
407#endif
408
409 /* Free the old one if we have one */
410 oldCancelConn = cancelConn;
411
412 /* be sure handle_sigint doesn't use pointer while freeing */
413 cancelConn = NULL;
414
415 if (oldCancelConn != NULL)
416 PQfreeCancel(oldCancelConn);
417
418 cancelConn = PQgetCancel(conn);
419
420#ifdef WIN32
421 LeaveCriticalSection(&cancelConnLock);
422#endif
423}
424
425/*
426 * ResetCancelConn
427 *
428 * Free the current cancel connection, if any, and set to NULL.
429 */
430void
431ResetCancelConn(void)
432{
433 PGcancel *oldCancelConn;
434
435#ifdef WIN32
436 EnterCriticalSection(&cancelConnLock);
437#endif
438
439 oldCancelConn = cancelConn;
440
441 /* be sure handle_sigint doesn't use pointer while freeing */
442 cancelConn = NULL;
443
444 if (oldCancelConn != NULL)
445 PQfreeCancel(oldCancelConn);
446
447#ifdef WIN32
448 LeaveCriticalSection(&cancelConnLock);
449#endif
450}
451
452#ifndef WIN32
453/*
454 * Handle interrupt signals by canceling the current command, if a cancelConn
455 * is set.
456 */
457static void
458handle_sigint(SIGNAL_ARGS)
459{
460 int save_errno = errno;
461 char errbuf[256];
462
463 /* Send QueryCancel if we are processing a database query */
464 if (cancelConn != NULL)
465 {
466 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
467 {
468 CancelRequested = true;
469 fprintf(stderr, _("Cancel request sent\n"));
470 }
471 else
472 fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
473 }
474 else
475 CancelRequested = true;
476
477 errno = save_errno; /* just in case the write changed it */
478}
479
480void
481setup_cancel_handler(void)
482{
483 pqsignal(SIGINT, handle_sigint);
484}
485#else /* WIN32 */
486
487/*
488 * Console control handler for Win32. Note that the control handler will
489 * execute on a *different thread* than the main one, so we need to do
490 * proper locking around those structures.
491 */
492static BOOL WINAPI
493consoleHandler(DWORD dwCtrlType)
494{
495 char errbuf[256];
496
497 if (dwCtrlType == CTRL_C_EVENT ||
498 dwCtrlType == CTRL_BREAK_EVENT)
499 {
500 /* Send QueryCancel if we are processing a database query */
501 EnterCriticalSection(&cancelConnLock);
502 if (cancelConn != NULL)
503 {
504 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
505 {
506 fprintf(stderr, _("Cancel request sent\n"));
507 CancelRequested = true;
508 }
509 else
510 fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
511 }
512 else
513 CancelRequested = true;
514
515 LeaveCriticalSection(&cancelConnLock);
516
517 return TRUE;
518 }
519 else
520 /* Return FALSE for any signals not being handled */
521 return FALSE;
522}
523
524void
525setup_cancel_handler(void)
526{
527 InitializeCriticalSection(&cancelConnLock);
528
529 SetConsoleCtrlHandler(consoleHandler, TRUE);
530}
531
532#endif /* WIN32 */
533