src/bin/scripts/common.c

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