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