granicus.if.org Git - postgresql/blob - src/bin/scripts/common.c

   1 /*-------------------------------------------------------------------------
   2  *
   3  *      common.c
   4  *              Common support routines for bin/scripts/
   5  *
   6  *
   7  * Portions Copyright (c) 1996-2018, 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
  22
  23 static PGcancel *volatile cancelConn = NULL;
  24 bool            CancelRequested = false;
  25
  26 #ifdef WIN32
  27 static CRITICAL_SECTION cancelConnLock;
  28 #endif
  29
  30 /*
  31  * Provide strictly harmonized handling of --help and --version
  32  * options.
  33  */
  34 void
  35 handle_help_version_opts(int argc, char *argv[],
  36                                                  const char *fixed_progname, help_handler hlp)
  37 {
  38         if (argc > 1)
  39         {
  40                 if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
  41                 {
  42                         hlp(get_progname(argv[0]));
  43                         exit(0);
  44                 }
  45                 if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
  46                 {
  47                         printf("%s (PostgreSQL) " PG_VERSION "\n", fixed_progname);
  48                         exit(0);
  49                 }
  50         }
  51 }
  52
  53
  54 /*
  55  * Make a database connection with the given parameters.
  56  *
  57  * An interactive password prompt is automatically issued if needed and
  58  * allowed by prompt_password.
  59  *
  60  * If allow_password_reuse is true, we will try to re-use any password
  61  * given during previous calls to this routine.  (Callers should not pass
  62  * allow_password_reuse=true unless reconnecting to the same database+user
  63  * as before, else we might create password exposure hazards.)
  64  */
  65 PGconn *
  66 connectDatabase(const char *dbname, const char *pghost, const char *pgport,
  67                                 const char *pguser, enum trivalue prompt_password,
  68                                 const char *progname, bool fail_ok, bool allow_password_reuse)
  69 {
  70         PGconn     *conn;
  71         bool            new_pass;
  72         static bool have_password = false;
  73         static char password[100];
  74
  75         if (!allow_password_reuse)
  76                 have_password = false;
  77
  78         if (!have_password && prompt_password == TRI_YES)
  79         {
  80                 simple_prompt("Password: ", password, sizeof(password), false);
  81                 have_password = true;
  82         }
  83
  84         /*
  85          * Start the connection.  Loop until we have a password if requested by
  86          * backend.
  87          */
  88         do
  89         {
  90                 const char *keywords[7];
  91                 const char *values[7];
  92
  93                 keywords[0] = "host";
  94                 values[0] = pghost;
  95                 keywords[1] = "port";
  96                 values[1] = pgport;
  97                 keywords[2] = "user";
  98                 values[2] = pguser;
  99                 keywords[3] = "password";
 100                 values[3] = have_password ? password : NULL;
 101                 keywords[4] = "dbname";
 102                 values[4] = dbname;
 103                 keywords[5] = "fallback_application_name";
 104                 values[5] = progname;
 105                 keywords[6] = NULL;
 106                 values[6] = NULL;
 107
 108                 new_pass = false;
 109                 conn = PQconnectdbParams(keywords, values, true);
 110
 111                 if (!conn)
 112                 {
 113                         fprintf(stderr, _("%s: could not connect to database %s: out of memory\n"),
 114                                         progname, dbname);
 115                         exit(1);
 116                 }
 117
 118                 /*
 119                  * No luck?  Trying asking (again) for a password.
 120                  */
 121                 if (PQstatus(conn) == CONNECTION_BAD &&
 122                         PQconnectionNeedsPassword(conn) &&
 123                         prompt_password != TRI_NO)
 124                 {
 125                         PQfinish(conn);
 126                         simple_prompt("Password: ", password, sizeof(password), false);
 127                         have_password = true;
 128                         new_pass = true;
 129                 }
 130         } while (new_pass);
 131
 132         /* check to see that the backend connection was successfully made */
 133         if (PQstatus(conn) == CONNECTION_BAD)
 134         {
 135                 if (fail_ok)
 136                 {
 137                         PQfinish(conn);
 138                         return NULL;
 139                 }
 140                 fprintf(stderr, _("%s: could not connect to database %s: %s"),
 141                                 progname, dbname, PQerrorMessage(conn));
 142                 exit(1);
 143         }
 144
 145         return conn;
 146 }
 147
 148 /*
 149  * Try to connect to the appropriate maintenance database.
 150  */
 151 PGconn *
 152 connectMaintenanceDatabase(const char *maintenance_db, const char *pghost,
 153                                                    const char *pgport, const char *pguser,
 154                                                    enum trivalue prompt_password,
 155                                                    const char *progname)
 156 {
 157         PGconn     *conn;
 158
 159         /* If a maintenance database name was specified, just connect to it. */
 160         if (maintenance_db)
 161                 return connectDatabase(maintenance_db, pghost, pgport, pguser,
 162                                                            prompt_password, progname, false, false);
 163
 164         /* Otherwise, try postgres first and then template1. */
 165         conn = connectDatabase("postgres", pghost, pgport, pguser, prompt_password,
 166                                                    progname, true, false);
 167         if (!conn)
 168                 conn = connectDatabase("template1", pghost, pgport, pguser,
 169                                                            prompt_password, progname, false, false);
 170
 171         return conn;
 172 }
 173
 174 /*
 175  * Run a query, return the results, exit program on failure.
 176  */
 177 PGresult *
 178 executeQuery(PGconn *conn, const char *query, const char *progname, bool echo)
 179 {
 180         PGresult   *res;
 181
 182         if (echo)
 183                 printf("%s\n", query);
 184
 185         res = PQexec(conn, query);
 186         if (!res ||
 187                 PQresultStatus(res) != PGRES_TUPLES_OK)
 188         {
 189                 fprintf(stderr, _("%s: query failed: %s"),
 190                                 progname, PQerrorMessage(conn));
 191                 fprintf(stderr, _("%s: query was: %s\n"),
 192                                 progname, query);
 193                 PQfinish(conn);
 194                 exit(1);
 195         }
 196
 197         return res;
 198 }
 199
 200
 201 /*
 202  * As above for a SQL command (which returns nothing).
 203  */
 204 void
 205 executeCommand(PGconn *conn, const char *query,
 206                            const char *progname, bool echo)
 207 {
 208         PGresult   *res;
 209
 210         if (echo)
 211                 printf("%s\n", query);
 212
 213         res = PQexec(conn, query);
 214         if (!res ||
 215                 PQresultStatus(res) != PGRES_COMMAND_OK)
 216         {
 217                 fprintf(stderr, _("%s: query failed: %s"),
 218                                 progname, PQerrorMessage(conn));
 219                 fprintf(stderr, _("%s: query was: %s\n"),
 220                                 progname, query);
 221                 PQfinish(conn);
 222                 exit(1);
 223         }
 224
 225         PQclear(res);
 226 }
 227
 228
 229 /*
 230  * As above for a SQL maintenance command (returns command success).
 231  * Command is executed with a cancel handler set, so Ctrl-C can
 232  * interrupt it.
 233  */
 234 bool
 235 executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
 236 {
 237         PGresult   *res;
 238         bool            r;
 239
 240         if (echo)
 241                 printf("%s\n", query);
 242
 243         SetCancelConn(conn);
 244         res = PQexec(conn, query);
 245         ResetCancelConn();
 246
 247         r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);
 248
 249         if (res)
 250                 PQclear(res);
 251
 252         return r;
 253 }
 254
 255 /*
 256  * Check yes/no answer in a localized way.  1=yes, 0=no, -1=neither.
 257  */
 258
 259 /* translator: abbreviation for "yes" */
 260 #define PG_YESLETTER gettext_noop("y")
 261 /* translator: abbreviation for "no" */
 262 #define PG_NOLETTER gettext_noop("n")
 263
 264 bool
 265 yesno_prompt(const char *question)
 266 {
 267         char            prompt[256];
 268
 269         /*------
 270            translator: This is a question followed by the translated options for
 271            "yes" and "no". */
 272         snprintf(prompt, sizeof(prompt), _("%s (%s/%s) "),
 273                          _(question), _(PG_YESLETTER), _(PG_NOLETTER));
 274
 275         for (;;)
 276         {
 277                 char            resp[10];
 278
 279                 simple_prompt(prompt, resp, sizeof(resp), true);
 280
 281                 if (strcmp(resp, _(PG_YESLETTER)) == 0)
 282                         return true;
 283                 if (strcmp(resp, _(PG_NOLETTER)) == 0)
 284                         return false;
 285
 286                 printf(_("Please answer \"%s\" or \"%s\".\n"),
 287                            _(PG_YESLETTER), _(PG_NOLETTER));
 288         }
 289 }
 290
 291 /*
 292  * SetCancelConn
 293  *
 294  * Set cancelConn to point to the current database connection.
 295  */
 296 void
 297 SetCancelConn(PGconn *conn)
 298 {
 299         PGcancel   *oldCancelConn;
 300
 301 #ifdef WIN32
 302         EnterCriticalSection(&cancelConnLock);
 303 #endif
 304
 305         /* Free the old one if we have one */
 306         oldCancelConn = cancelConn;
 307
 308         /* be sure handle_sigint doesn't use pointer while freeing */
 309         cancelConn = NULL;
 310
 311         if (oldCancelConn != NULL)
 312                 PQfreeCancel(oldCancelConn);
 313
 314         cancelConn = PQgetCancel(conn);
 315
 316 #ifdef WIN32
 317         LeaveCriticalSection(&cancelConnLock);
 318 #endif
 319 }
 320
 321 /*
 322  * ResetCancelConn
 323  *
 324  * Free the current cancel connection, if any, and set to NULL.
 325  */
 326 void
 327 ResetCancelConn(void)
 328 {
 329         PGcancel   *oldCancelConn;
 330
 331 #ifdef WIN32
 332         EnterCriticalSection(&cancelConnLock);
 333 #endif
 334
 335         oldCancelConn = cancelConn;
 336
 337         /* be sure handle_sigint doesn't use pointer while freeing */
 338         cancelConn = NULL;
 339
 340         if (oldCancelConn != NULL)
 341                 PQfreeCancel(oldCancelConn);
 342
 343 #ifdef WIN32
 344         LeaveCriticalSection(&cancelConnLock);
 345 #endif
 346 }
 347
 348 #ifndef WIN32
 349 /*
 350  * Handle interrupt signals by canceling the current command, if a cancelConn
 351  * is set.
 352  */
 353 static void
 354 handle_sigint(SIGNAL_ARGS)
 355 {
 356         int                     save_errno = errno;
 357         char            errbuf[256];
 358
 359         /* Send QueryCancel if we are processing a database query */
 360         if (cancelConn != NULL)
 361         {
 362                 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 363                 {
 364                         CancelRequested = true;
 365                         fprintf(stderr, _("Cancel request sent\n"));
 366                 }
 367                 else
 368                         fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
 369         }
 370         else
 371                 CancelRequested = true;
 372
 373         errno = save_errno;                     /* just in case the write changed it */
 374 }
 375
 376 void
 377 setup_cancel_handler(void)
 378 {
 379         pqsignal(SIGINT, handle_sigint);
 380 }
 381 #else                                                   /* WIN32 */
 382
 383 /*
 384  * Console control handler for Win32. Note that the control handler will
 385  * execute on a *different thread* than the main one, so we need to do
 386  * proper locking around those structures.
 387  */
 388 static BOOL WINAPI
 389 consoleHandler(DWORD dwCtrlType)
 390 {
 391         char            errbuf[256];
 392
 393         if (dwCtrlType == CTRL_C_EVENT ||
 394                 dwCtrlType == CTRL_BREAK_EVENT)
 395         {
 396                 /* Send QueryCancel if we are processing a database query */
 397                 EnterCriticalSection(&cancelConnLock);
 398                 if (cancelConn != NULL)
 399                 {
 400                         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 401                         {
 402                                 fprintf(stderr, _("Cancel request sent\n"));
 403                                 CancelRequested = true;
 404                         }
 405                         else
 406                                 fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
 407                 }
 408                 else
 409                         CancelRequested = true;
 410
 411                 LeaveCriticalSection(&cancelConnLock);
 412
 413                 return TRUE;
 414         }
 415         else
 416                 /* Return FALSE for any signals not being handled */
 417                 return FALSE;
 418 }
 419
 420 void
 421 setup_cancel_handler(void)
 422 {
 423         InitializeCriticalSection(&cancelConnLock);
 424
 425         SetConsoleCtrlHandler(consoleHandler, TRUE);
 426 }
 427
 428 #endif                                                  /* WIN32 */