src/bin/psql/common.c

   1 /*
   2  * psql - the PostgreSQL interactive terminal
   3  *
   4  * Copyright (c) 2000-2005, PostgreSQL Global Development Group
   5  *
   6  * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.108 2005/10/15 02:49:40 momjian Exp $
   7  */
   8 #include "postgres_fe.h"
   9 #include "common.h"
  10
  11 #include <ctype.h>
  12 #ifndef HAVE_STRDUP
  13 #include <strdup.h>
  14 #endif
  15 #include <signal.h>
  16 #ifndef WIN32
  17 #include <sys/time.h>
  18 #include <unistd.h>                             /* for write() */
  19 #include <setjmp.h>
  20 #else
  21 #include <io.h>                                 /* for _write() */
  22 #include <win32.h>
  23 #include <sys/timeb.h>                  /* for _ftime() */
  24 #endif
  25
  26 #include "libpq-fe.h"
  27 #include "pqsignal.h"
  28
  29 #include "settings.h"
  30 #include "variables.h"
  31 #include "command.h"
  32 #include "copy.h"
  33 #include "prompt.h"
  34 #include "print.h"
  35 #include "mainloop.h"
  36 #include "mb/pg_wchar.h"
  37
  38
  39 /* Workarounds for Windows */
  40 /* Probably to be moved up the source tree in the future, perhaps to be replaced by
  41  * more specific checks like configure-style HAVE_GETTIMEOFDAY macros.
  42  */
  43 #ifndef WIN32
  44
  45 typedef struct timeval TimevalStruct;
  46
  47 #define GETTIMEOFDAY(T) gettimeofday(T, NULL)
  48 #define DIFF_MSEC(T, U) \
  49         ((((int) ((T)->tv_sec - (U)->tv_sec)) * 1000000.0 + \
  50           ((int) ((T)->tv_usec - (U)->tv_usec))) / 1000.0)
  51 #else
  52
  53 typedef struct _timeb TimevalStruct;
  54
  55 #define GETTIMEOFDAY(T) _ftime(T)
  56 #define DIFF_MSEC(T, U) \
  57         (((T)->time - (U)->time) * 1000.0 + \
  58          ((T)->millitm - (U)->millitm))
  59 #endif
  60
  61 extern bool prompt_state;
  62
  63
  64 static bool command_no_begin(const char *query);
  65
  66 /*
  67  * "Safe" wrapper around strdup()
  68  */
  69 char *
  70 pg_strdup(const char *string)
  71 {
  72         char       *tmp;
  73
  74         if (!string)
  75         {
  76                 fprintf(stderr, _("%s: pg_strdup: cannot duplicate null pointer (internal error)\n"),
  77                                 pset.progname);
  78                 exit(EXIT_FAILURE);
  79         }
  80         tmp = strdup(string);
  81         if (!tmp)
  82         {
  83                 psql_error("out of memory\n");
  84                 exit(EXIT_FAILURE);
  85         }
  86         return tmp;
  87 }
  88
  89 void *
  90 pg_malloc(size_t size)
  91 {
  92         void       *tmp;
  93
  94         tmp = malloc(size);
  95         if (!tmp)
  96         {
  97                 psql_error("out of memory\n");
  98                 exit(EXIT_FAILURE);
  99         }
 100         return tmp;
 101 }
 102
 103 void *
 104 pg_malloc_zero(size_t size)
 105 {
 106         void       *tmp;
 107
 108         tmp = pg_malloc(size);
 109         memset(tmp, 0, size);
 110         return tmp;
 111 }
 112
 113 void *
 114 pg_calloc(size_t nmemb, size_t size)
 115 {
 116         void       *tmp;
 117
 118         tmp = calloc(nmemb, size);
 119         if (!tmp)
 120         {
 121                 psql_error("out of memory");
 122                 exit(EXIT_FAILURE);
 123         }
 124         return tmp;
 125 }
 126
 127 /*
 128  * setQFout
 129  * -- handler for -o command line option and \o command
 130  *
 131  * Tries to open file fname (or pipe if fname starts with '|')
 132  * and stores the file handle in pset)
 133  * Upon failure, sets stdout and returns false.
 134  */
 135 bool
 136 setQFout(const char *fname)
 137 {
 138         bool            status = true;
 139
 140         /* Close old file/pipe */
 141         if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
 142         {
 143                 if (pset.queryFoutPipe)
 144                         pclose(pset.queryFout);
 145                 else
 146                         fclose(pset.queryFout);
 147         }
 148
 149         /* If no filename, set stdout */
 150         if (!fname || fname[0] == '\0')
 151         {
 152                 pset.queryFout = stdout;
 153                 pset.queryFoutPipe = false;
 154         }
 155         else if (*fname == '|')
 156         {
 157                 pset.queryFout = popen(fname + 1, "w");
 158                 pset.queryFoutPipe = true;
 159         }
 160         else
 161         {
 162                 pset.queryFout = fopen(fname, "w");
 163                 pset.queryFoutPipe = false;
 164         }
 165
 166         if (!(pset.queryFout))
 167         {
 168                 psql_error("%s: %s\n", fname, strerror(errno));
 169                 pset.queryFout = stdout;
 170                 pset.queryFoutPipe = false;
 171                 status = false;
 172         }
 173
 174         /* Direct signals */
 175 #ifndef WIN32
 176         pqsignal(SIGPIPE, pset.queryFoutPipe ? SIG_IGN : SIG_DFL);
 177 #endif
 178
 179         return status;
 180 }
 181
 182
 183
 184 /*
 185  * Error reporting for scripts. Errors should look like
 186  *       psql:filename:lineno: message
 187  *
 188  */
 189 void
 190 psql_error(const char *fmt,...)
 191 {
 192         va_list         ap;
 193
 194         fflush(stdout);
 195         if (pset.queryFout != stdout)
 196                 fflush(pset.queryFout);
 197
 198         if (pset.inputfile)
 199                 fprintf(stderr, "%s:%s:%u: ", pset.progname, pset.inputfile, pset.lineno);
 200         va_start(ap, fmt);
 201         vfprintf(stderr, _(fmt), ap);
 202         va_end(ap);
 203 }
 204
 205
 206
 207 /*
 208  * for backend Notice messages (INFO, WARNING, etc)
 209  */
 210 void
 211 NoticeProcessor(void *arg, const char *message)
 212 {
 213         (void) arg;                                     /* not used */
 214         psql_error("%s", message);
 215 }
 216
 217
 218
 219 /*
 220  * Code to support query cancellation
 221  *
 222  * Before we start a query, we enable a SIGINT signal catcher that sends a
 223  * cancel request to the backend. Note that sending the cancel directly from
 224  * the signal handler is safe because PQcancel() is written to make it
 225  * so. We use write() to print to stderr because it's better to use simple
 226  * facilities in a signal handler.
 227  *
 228  * On win32, the signal cancelling happens on a separate thread, because
 229  * that's how SetConsoleCtrlHandler works. The PQcancel function is safe
 230  * for this (unlike PQrequestCancel). However, a CRITICAL_SECTION is required
 231  * to protect the PGcancel structure against being changed while the other
 232  * thread is using it.
 233  */
 234 static PGcancel *cancelConn = NULL;
 235
 236 #ifdef WIN32
 237 static CRITICAL_SECTION cancelConnLock;
 238 #endif
 239
 240 volatile bool cancel_pressed = false;
 241
 242 #define write_stderr(str)       write(fileno(stderr), str, strlen(str))
 243
 244
 245 #ifndef WIN32
 246
 247 void
 248 handle_sigint(SIGNAL_ARGS)
 249 {
 250         int                     save_errno = errno;
 251         char            errbuf[256];
 252
 253         /* Don't muck around if prompting for a password. */
 254         if (prompt_state)
 255                 return;
 256
 257         if (cancelConn == NULL)
 258                 siglongjmp(main_loop_jmp, 1);
 259
 260         cancel_pressed = true;
 261
 262         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 263                 write_stderr("Cancel request sent\n");
 264         else
 265         {
 266                 write_stderr("Could not send cancel request: ");
 267                 write_stderr(errbuf);
 268         }
 269         errno = save_errno;                     /* just in case the write changed it */
 270 }
 271 #else                                                   /* WIN32 */
 272
 273 static BOOL WINAPI
 274 consoleHandler(DWORD dwCtrlType)
 275 {
 276         char            errbuf[256];
 277
 278         if (dwCtrlType == CTRL_C_EVENT ||
 279                 dwCtrlType == CTRL_BREAK_EVENT)
 280         {
 281                 if (prompt_state)
 282                         return TRUE;
 283
 284                 /* Perform query cancel */
 285                 EnterCriticalSection(&cancelConnLock);
 286                 if (cancelConn != NULL)
 287                 {
 288                         cancel_pressed = true;
 289
 290                         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 291                                 write_stderr("Cancel request sent\n");
 292                         else
 293                         {
 294                                 write_stderr("Could not send cancel request: ");
 295                                 write_stderr(errbuf);
 296                         }
 297                 }
 298                 LeaveCriticalSection(&cancelConnLock);
 299
 300                 return TRUE;
 301         }
 302         else
 303                 /* Return FALSE for any signals not being handled */
 304                 return FALSE;
 305 }
 306
 307 void
 308 setup_win32_locks(void)
 309 {
 310         InitializeCriticalSection(&cancelConnLock);
 311 }
 312
 313 void
 314 setup_cancel_handler(void)
 315 {
 316         SetConsoleCtrlHandler(consoleHandler, TRUE);
 317 }
 318 #endif   /* WIN32 */
 319
 320
 321 /* ConnectionUp
 322  *
 323  * Returns whether our backend connection is still there.
 324  */
 325 static bool
 326 ConnectionUp(void)
 327 {
 328         return PQstatus(pset.db) != CONNECTION_BAD;
 329 }
 330
 331
 332
 333 /* CheckConnection
 334  *
 335  * Verify that we still have a good connection to the backend, and if not,
 336  * see if it can be restored.
 337  *
 338  * Returns true if either the connection was still there, or it could be
 339  * restored successfully; false otherwise.      If, however, there was no
 340  * connection and the session is non-interactive, this will exit the program
 341  * with a code of EXIT_BADCONN.
 342  */
 343 static bool
 344 CheckConnection(void)
 345 {
 346         bool            OK;
 347
 348         OK = ConnectionUp();
 349         if (!OK)
 350         {
 351                 if (!pset.cur_cmd_interactive)
 352                 {
 353                         psql_error("connection to server was lost\n");
 354                         exit(EXIT_BADCONN);
 355                 }
 356
 357                 fputs(_("The connection to the server was lost. Attempting reset: "), stderr);
 358                 PQreset(pset.db);
 359                 OK = ConnectionUp();
 360                 if (!OK)
 361                 {
 362                         fputs(_("Failed.\n"), stderr);
 363                         PQfinish(pset.db);
 364                         pset.db = NULL;
 365                         ResetCancelConn();
 366                         UnsyncVariables();
 367                 }
 368                 else
 369                         fputs(_("Succeeded.\n"), stderr);
 370         }
 371
 372         return OK;
 373 }
 374
 375
 376
 377 /*
 378  * SetCancelConn
 379  *
 380  * Set cancelConn to point to the current database connection.
 381  */
 382 static void
 383 SetCancelConn(void)
 384 {
 385 #ifdef WIN32
 386         EnterCriticalSection(&cancelConnLock);
 387 #endif
 388
 389         /* Free the old one if we have one */
 390         if (cancelConn != NULL)
 391                 PQfreeCancel(cancelConn);
 392
 393         cancelConn = PQgetCancel(pset.db);
 394
 395 #ifdef WIN32
 396         LeaveCriticalSection(&cancelConnLock);
 397 #endif
 398 }
 399
 400
 401 /*
 402  * ResetCancelConn
 403  *
 404  * Free the current cancel connection, if any, and set to NULL.
 405  */
 406 void
 407 ResetCancelConn(void)
 408 {
 409 #ifdef WIN32
 410         EnterCriticalSection(&cancelConnLock);
 411 #endif
 412
 413         if (cancelConn)
 414                 PQfreeCancel(cancelConn);
 415
 416         cancelConn = NULL;
 417
 418 #ifdef WIN32
 419         LeaveCriticalSection(&cancelConnLock);
 420 #endif
 421 }
 422
 423
 424 /*
 425  * on errors, print syntax error position if available.
 426  *
 427  * the query is expected to be in the client encoding.
 428  */
 429 static void
 430 ReportSyntaxErrorPosition(const PGresult *result, const char *query)
 431 {
 432 #define DISPLAY_SIZE    60              /* screen width limit, in screen cols */
 433 #define MIN_RIGHT_CUT   10              /* try to keep this far away from EOL */
 434
 435         int                     loc = 0;
 436         const char *sp;
 437         int                     clen,
 438                                 slen,
 439                                 i,
 440                            *qidx,
 441                            *scridx,
 442                                 qoffset,
 443                                 scroffset,
 444                                 ibeg,
 445                                 iend,
 446                                 loc_line;
 447         char       *wquery;
 448         bool            beg_trunc,
 449                                 end_trunc;
 450         PQExpBufferData msg;
 451
 452         if (pset.verbosity == PQERRORS_TERSE)
 453                 return;
 454
 455         sp = PQresultErrorField(result, PG_DIAG_STATEMENT_POSITION);
 456         if (sp == NULL)
 457         {
 458                 sp = PQresultErrorField(result, PG_DIAG_INTERNAL_POSITION);
 459                 if (sp == NULL)
 460                         return;                         /* no syntax error */
 461                 query = PQresultErrorField(result, PG_DIAG_INTERNAL_QUERY);
 462         }
 463         if (query == NULL)
 464                 return;                                 /* nothing to reference location to */
 465
 466         if (sscanf(sp, "%d", &loc) != 1)
 467         {
 468                 psql_error("INTERNAL ERROR: unexpected statement position \"%s\"\n",
 469                                    sp);
 470                 return;
 471         }
 472
 473         /* Make a writable copy of the query, and a buffer for messages. */
 474         wquery = pg_strdup(query);
 475
 476         initPQExpBuffer(&msg);
 477
 478         /*
 479          * The returned cursor position is measured in logical characters. Each
 480          * character might occupy multiple physical bytes in the string, and in
 481          * some Far Eastern character sets it might take more than one screen
 482          * column as well.      We compute the starting byte offset and starting
 483          * screen column of each logical character, and store these in qidx[] and
 484          * scridx[] respectively.
 485          */
 486
 487         /* we need a safe allocation size... */
 488         slen = strlen(query) + 1;
 489
 490         qidx = (int *) pg_malloc(slen * sizeof(int));
 491         scridx = (int *) pg_malloc(slen * sizeof(int));
 492
 493         qoffset = 0;
 494         scroffset = 0;
 495         for (i = 0; query[qoffset] != '\0'; i++)
 496         {
 497                 qidx[i] = qoffset;
 498                 scridx[i] = scroffset;
 499                 scroffset += PQdsplen(&query[qoffset], pset.encoding);
 500                 qoffset += PQmblen(&query[qoffset], pset.encoding);
 501         }
 502         qidx[i] = qoffset;
 503         scridx[i] = scroffset;
 504         clen = i;
 505         psql_assert(clen < slen);
 506
 507         /* convert loc to zero-based offset in qidx/scridx arrays */
 508         loc--;
 509
 510         /* do we have something to show? */
 511         if (loc >= 0 && loc <= clen)
 512         {
 513                 /* input line number of our syntax error. */
 514                 loc_line = 1;
 515                 /* first included char of extract. */
 516                 ibeg = 0;
 517                 /* last-plus-1 included char of extract. */
 518                 iend = clen;
 519
 520                 /*
 521                  * Replace tabs with spaces in the writable copy.  (Later we might
 522                  * want to think about coping with their variable screen width, but
 523                  * not today.)
 524                  *
 525                  * Extract line number and begin and end indexes of line containing error
 526                  * location.  There will not be any newlines or carriage returns in
 527                  * the selected extract.
 528                  */
 529                 for (i = 0; i < clen; i++)
 530                 {
 531                         /* character length must be 1 or it's not ASCII */
 532                         if ((qidx[i + 1] - qidx[i]) == 1)
 533                         {
 534                                 if (wquery[qidx[i]] == '\t')
 535                                         wquery[qidx[i]] = ' ';
 536                                 else if (wquery[qidx[i]] == '\r' || wquery[qidx[i]] == '\n')
 537                                 {
 538                                         if (i < loc)
 539                                         {
 540                                                 /*
 541                                                  * count lines before loc.      Each \r or \n counts as a
 542                                                  * line except when \r \n appear together.
 543                                                  */
 544                                                 if (wquery[qidx[i]] == '\r' ||
 545                                                         i == 0 ||
 546                                                         (qidx[i] - qidx[i - 1]) != 1 ||
 547                                                         wquery[qidx[i - 1]] != '\r')
 548                                                         loc_line++;
 549                                                 /* extract beginning = last line start before loc. */
 550                                                 ibeg = i + 1;
 551                                         }
 552                                         else
 553                                         {
 554                                                 /* set extract end. */
 555                                                 iend = i;
 556                                                 /* done scanning. */
 557                                                 break;
 558                                         }
 559                                 }
 560                         }
 561                 }
 562
 563                 /* If the line extracted is too long, we truncate it. */
 564                 beg_trunc = false;
 565                 end_trunc = false;
 566                 if (scridx[iend] - scridx[ibeg] > DISPLAY_SIZE)
 567                 {
 568                         /*
 569                          * We first truncate right if it is enough.  This code might be
 570                          * off a space or so on enforcing MIN_RIGHT_CUT if there's a wide
 571                          * character right there, but that should be okay.
 572                          */
 573                         if (scridx[ibeg] + DISPLAY_SIZE >= scridx[loc] + MIN_RIGHT_CUT)
 574                         {
 575                                 while (scridx[iend] - scridx[ibeg] > DISPLAY_SIZE)
 576                                         iend--;
 577                                 end_trunc = true;
 578                         }
 579                         else
 580                         {
 581                                 /* Truncate right if not too close to loc. */
 582                                 while (scridx[loc] + MIN_RIGHT_CUT < scridx[iend])
 583                                 {
 584                                         iend--;
 585                                         end_trunc = true;
 586                                 }
 587
 588                                 /* Truncate left if still too long. */
 589                                 while (scridx[iend] - scridx[ibeg] > DISPLAY_SIZE)
 590                                 {
 591                                         ibeg++;
 592                                         beg_trunc = true;
 593                                 }
 594                         }
 595                 }
 596
 597                 /* the extract MUST contain the target position! */
 598                 psql_assert(ibeg <= loc && loc <= iend);
 599
 600                 /* truncate working copy at desired endpoint */
 601                 wquery[qidx[iend]] = '\0';
 602
 603                 /* Begin building the finished message. */
 604                 printfPQExpBuffer(&msg, _("LINE %d: "), loc_line);
 605                 if (beg_trunc)
 606                         appendPQExpBufferStr(&msg, "...");
 607
 608                 /*
 609                  * While we have the prefix in the msg buffer, compute its screen
 610                  * width.
 611                  */
 612                 scroffset = 0;
 613                 for (i = 0; i < msg.len; i += PQmblen(&msg.data[i], pset.encoding))
 614                         scroffset += PQdsplen(&msg.data[i], pset.encoding);
 615
 616                 /* Finish and emit the message. */
 617                 appendPQExpBufferStr(&msg, &wquery[qidx[ibeg]]);
 618                 if (end_trunc)
 619                         appendPQExpBufferStr(&msg, "...");
 620
 621                 psql_error("%s\n", msg.data);
 622
 623                 /* Now emit the cursor marker line. */
 624                 scroffset += scridx[loc] - scridx[ibeg];
 625                 resetPQExpBuffer(&msg);
 626                 for (i = 0; i < scroffset; i++)
 627                         appendPQExpBufferChar(&msg, ' ');
 628                 appendPQExpBufferChar(&msg, '^');
 629
 630                 psql_error("%s\n", msg.data);
 631         }
 632
 633         /* Clean up. */
 634         termPQExpBuffer(&msg);
 635
 636         free(wquery);
 637         free(qidx);
 638         free(scridx);
 639 }
 640
 641
 642 /*
 643  * AcceptResult
 644  *
 645  * Checks whether a result is valid, giving an error message if necessary;
 646  * resets cancelConn as needed, and ensures that the connection to the backend
 647  * is still up.
 648  *
 649  * Returns true for valid result, false for error state.
 650  */
 651 static bool
 652 AcceptResult(const PGresult *result, const char *query)
 653 {
 654         bool            OK = true;
 655
 656         ResetCancelConn();
 657
 658         if (!result)
 659                 OK = false;
 660         else
 661                 switch (PQresultStatus(result))
 662                 {
 663                         case PGRES_COMMAND_OK:
 664                         case PGRES_TUPLES_OK:
 665                         case PGRES_EMPTY_QUERY:
 666                         case PGRES_COPY_IN:
 667                                 /* Fine, do nothing */
 668                                 break;
 669
 670                         case PGRES_COPY_OUT:
 671                                 /* keep cancel connection for copy out state */
 672                                 SetCancelConn();
 673                                 break;
 674
 675                         default:
 676                                 OK = false;
 677                                 break;
 678                 }
 679
 680         if (!OK)
 681         {
 682                 const char *error = PQerrorMessage(pset.db);
 683
 684                 if (strlen(error))
 685                         psql_error("%s", error);
 686
 687                 ReportSyntaxErrorPosition(result, query);
 688                 CheckConnection();
 689         }
 690
 691         return OK;
 692 }
 693
 694
 695
 696 /*
 697  * PSQLexec
 698  *
 699  * This is the way to send "backdoor" queries (those not directly entered
 700  * by the user). It is subject to -E but not -e.
 701  *
 702  * In autocommit-off mode, a new transaction block is started if start_xact
 703  * is true; nothing special is done when start_xact is false.  Typically,
 704  * start_xact = false is used for SELECTs and explicit BEGIN/COMMIT commands.
 705  *
 706  * Note: we don't bother to check PQclientEncoding; it is assumed that no
 707  * caller uses this path to issue "SET CLIENT_ENCODING".
 708  */
 709 PGresult *
 710 PSQLexec(const char *query, bool start_xact)
 711 {
 712         PGresult   *res;
 713         int                     echo_hidden;
 714
 715         if (!pset.db)
 716         {
 717                 psql_error("You are currently not connected to a database.\n");
 718                 return NULL;
 719         }
 720
 721         echo_hidden = SwitchVariable(pset.vars, "ECHO_HIDDEN", "noexec", NULL);
 722         if (echo_hidden != VAR_NOTSET)
 723         {
 724                 printf(_("********* QUERY **********\n"
 725                                  "%s\n"
 726                                  "**************************\n\n"), query);
 727                 fflush(stdout);
 728                 if (pset.logfile)
 729                 {
 730                         fprintf(pset.logfile,
 731                                         _("********* QUERY **********\n"
 732                                           "%s\n"
 733                                           "**************************\n\n"), query);
 734                         fflush(pset.logfile);
 735                 }
 736
 737                 if (echo_hidden == 1)   /* noexec? */
 738                         return NULL;
 739         }
 740
 741         SetCancelConn();
 742
 743         if (start_xact && PQtransactionStatus(pset.db) == PQTRANS_IDLE &&
 744                 !GetVariableBool(pset.vars, "AUTOCOMMIT"))
 745         {
 746                 res = PQexec(pset.db, "BEGIN");
 747                 if (PQresultStatus(res) != PGRES_COMMAND_OK)
 748                 {
 749                         psql_error("%s", PQerrorMessage(pset.db));
 750                         PQclear(res);
 751                         ResetCancelConn();
 752                         return NULL;
 753                 }
 754                 PQclear(res);
 755         }
 756
 757         res = PQexec(pset.db, query);
 758
 759         if (!AcceptResult(res, query) && res)
 760         {
 761                 PQclear(res);
 762                 res = NULL;
 763         }
 764
 765         return res;
 766 }
 767
 768
 769
 770 /*
 771  * PrintNotifications: check for asynchronous notifications, and print them out
 772  */
 773 static void
 774 PrintNotifications(void)
 775 {
 776         PGnotify   *notify;
 777
 778         while ((notify = PQnotifies(pset.db)))
 779         {
 780                 fprintf(pset.queryFout, _("Asynchronous notification \"%s\" received from server process with PID %d.\n"),
 781                                 notify->relname, notify->be_pid);
 782                 fflush(pset.queryFout);
 783                 PQfreemem(notify);
 784         }
 785 }
 786
 787
 788 /*
 789  * PrintQueryTuples: assuming query result is OK, print its tuples
 790  *
 791  * Returns true if successful, false otherwise.
 792  */
 793 static bool
 794 PrintQueryTuples(const PGresult *results)
 795 {
 796         printQueryOpt my_popt = pset.popt;
 797
 798         my_popt.topt.normal_query = true;
 799
 800         /* write output to \g argument, if any */
 801         if (pset.gfname)
 802         {
 803                 FILE       *queryFout_copy = pset.queryFout;
 804                 bool            queryFoutPipe_copy = pset.queryFoutPipe;
 805
 806                 pset.queryFout = stdout;        /* so it doesn't get closed */
 807
 808                 /* open file/pipe */
 809                 if (!setQFout(pset.gfname))
 810                 {
 811                         pset.queryFout = queryFout_copy;
 812                         pset.queryFoutPipe = queryFoutPipe_copy;
 813                         return false;
 814                 }
 815
 816                 printQuery(results, &my_popt, pset.queryFout, pset.logfile);
 817
 818                 /* close file/pipe, restore old setting */
 819                 setQFout(NULL);
 820
 821                 pset.queryFout = queryFout_copy;
 822                 pset.queryFoutPipe = queryFoutPipe_copy;
 823
 824                 free(pset.gfname);
 825                 pset.gfname = NULL;
 826         }
 827         else
 828                 printQuery(results, &my_popt, pset.queryFout, pset.logfile);
 829
 830         return true;
 831 }
 832
 833
 834 /*
 835  * ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it
 836  *
 837  * Note: Utility function for use by SendQuery() only.
 838  *
 839  * Returns true if the query executed successfully, false otherwise.
 840  */
 841 static bool
 842 ProcessCopyResult(PGresult *results)
 843 {
 844         bool            success = false;
 845
 846         if (!results)
 847                 return false;
 848
 849         switch (PQresultStatus(results))
 850         {
 851                 case PGRES_TUPLES_OK:
 852                 case PGRES_COMMAND_OK:
 853                 case PGRES_EMPTY_QUERY:
 854                         /* nothing to do here */
 855                         success = true;
 856                         break;
 857
 858                 case PGRES_COPY_OUT:
 859                         success = handleCopyOut(pset.db, pset.queryFout);
 860                         break;
 861
 862                 case PGRES_COPY_IN:
 863                         success = handleCopyIn(pset.db, pset.cur_cmd_source);
 864                         break;
 865
 866                 default:
 867                         break;
 868         }
 869
 870         /* may need this to recover from conn loss during COPY */
 871         if (!CheckConnection())
 872                 return false;
 873
 874         return success;
 875 }
 876
 877
 878 /*
 879  * PrintQueryResults: print out query results as required
 880  *
 881  * Note: Utility function for use by SendQuery() only.
 882  *
 883  * Returns true if the query executed successfully, false otherwise.
 884  */
 885 static bool
 886 PrintQueryResults(PGresult *results)
 887 {
 888         bool            success = false;
 889
 890         if (!results)
 891                 return false;
 892
 893         switch (PQresultStatus(results))
 894         {
 895                 case PGRES_TUPLES_OK:
 896                         success = PrintQueryTuples(results);
 897                         break;
 898
 899                 case PGRES_COMMAND_OK:
 900                         {
 901                                 char            buf[10];
 902
 903                                 success = true;
 904                                 snprintf(buf, sizeof(buf),
 905                                                  "%u", (unsigned int) PQoidValue(results));
 906                                 if (!QUIET())
 907                                 {
 908                                         if (pset.popt.topt.format == PRINT_HTML)
 909                                         {
 910                                                 fputs("<p>", pset.queryFout);
 911                                                 html_escaped_print(PQcmdStatus(results),
 912                                                                                    pset.queryFout);
 913                                                 fputs("</p>\n", pset.queryFout);
 914                                         }
 915                                         else
 916                                                 fprintf(pset.queryFout, "%s\n", PQcmdStatus(results));
 917                                 }
 918                                 if (pset.logfile)
 919                                         fprintf(pset.logfile, "%s\n", PQcmdStatus(results));
 920                                 SetVariable(pset.vars, "LASTOID", buf);
 921                                 break;
 922                         }
 923
 924                 case PGRES_EMPTY_QUERY:
 925                         success = true;
 926                         break;
 927
 928                 case PGRES_COPY_OUT:
 929                 case PGRES_COPY_IN:
 930                         /* nothing to do here */
 931                         success = true;
 932                         break;
 933
 934                 default:
 935                         break;
 936         }
 937
 938         fflush(pset.queryFout);
 939
 940         return success;
 941 }
 942
 943
 944 /*
 945  * SendQuery: send the query string to the backend
 946  * (and print out results)
 947  *
 948  * Note: This is the "front door" way to send a query. That is, use it to
 949  * send queries actually entered by the user. These queries will be subject to
 950  * single step mode.
 951  * To send "back door" queries (generated by slash commands, etc.) in a
 952  * controlled way, use PSQLexec().
 953  *
 954  * Returns true if the query executed successfully, false otherwise.
 955  */
 956 bool
 957 SendQuery(const char *query)
 958 {
 959         PGresult   *results;
 960         TimevalStruct before,
 961                                 after;
 962         bool            OK,
 963                                 on_error_rollback_savepoint = false;
 964         PGTransactionStatusType transaction_status;
 965         static bool on_error_rollback_warning = false;
 966         const char *rollback_str;
 967
 968         if (!pset.db)
 969         {
 970                 psql_error("You are currently not connected to a database.\n");
 971                 return false;
 972         }
 973
 974         if (GetVariableBool(pset.vars, "SINGLESTEP"))
 975         {
 976                 char            buf[3];
 977
 978                 printf(_("***(Single step mode: verify command)*******************************************\n"
 979                                  "%s\n"
 980                                  "***(press return to proceed or enter x and return to cancel)********************\n"),
 981                            query);
 982                 fflush(stdout);
 983                 if (fgets(buf, sizeof(buf), stdin) != NULL)
 984                         if (buf[0] == 'x')
 985                                 return false;
 986         }
 987         else if (VariableEquals(pset.vars, "ECHO", "queries"))
 988         {
 989                 puts(query);
 990                 fflush(stdout);
 991         }
 992
 993         if (pset.logfile)
 994         {
 995                 fprintf(pset.logfile,
 996                                 _("********* QUERY **********\n"
 997                                   "%s\n"
 998                                   "**************************\n\n"), query);
 999                 fflush(pset.logfile);
1000         }
1001
1002         SetCancelConn();
1003
1004         transaction_status = PQtransactionStatus(pset.db);
1005
1006         if (transaction_status == PQTRANS_IDLE &&
1007                 !GetVariableBool(pset.vars, "AUTOCOMMIT") &&
1008                 !command_no_begin(query))
1009         {
1010                 results = PQexec(pset.db, "BEGIN");
1011                 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1012                 {
1013                         psql_error("%s", PQerrorMessage(pset.db));
1014                         PQclear(results);
1015                         ResetCancelConn();
1016                         return false;
1017                 }
1018                 PQclear(results);
1019                 transaction_status = PQtransactionStatus(pset.db);
1020         }
1021
1022         if (transaction_status == PQTRANS_INTRANS &&
1023           (rollback_str = GetVariable(pset.vars, "ON_ERROR_ROLLBACK")) != NULL &&
1024         /* !off and !interactive is 'on' */
1025                 pg_strcasecmp(rollback_str, "off") != 0 &&
1026                 (pset.cur_cmd_interactive ||
1027                  pg_strcasecmp(rollback_str, "interactive") != 0))
1028         {
1029                 if (on_error_rollback_warning == false && pset.sversion < 80000)
1030                 {
1031                         fprintf(stderr, _("The server version (%d) does not support savepoints for ON_ERROR_ROLLBACK.\n"),
1032                                         pset.sversion);
1033                         on_error_rollback_warning = true;
1034                 }
1035                 else
1036                 {
1037                         results = PQexec(pset.db, "SAVEPOINT pg_psql_temporary_savepoint");
1038                         if (PQresultStatus(results) != PGRES_COMMAND_OK)
1039                         {
1040                                 psql_error("%s", PQerrorMessage(pset.db));
1041                                 PQclear(results);
1042                                 ResetCancelConn();
1043                                 return false;
1044                         }
1045                         PQclear(results);
1046                         on_error_rollback_savepoint = true;
1047                 }
1048         }
1049
1050         if (pset.timing)
1051                 GETTIMEOFDAY(&before);
1052
1053         results = PQexec(pset.db, query);
1054
1055         /* these operations are included in the timing result: */
1056         OK = (AcceptResult(results, query) && ProcessCopyResult(results));
1057
1058         if (pset.timing)
1059                 GETTIMEOFDAY(&after);
1060
1061         /* but printing results isn't: */
1062         if (OK)
1063                 OK = PrintQueryResults(results);
1064
1065         PQclear(results);
1066
1067         /* If we made a temporary savepoint, possibly release/rollback */
1068         if (on_error_rollback_savepoint)
1069         {
1070                 transaction_status = PQtransactionStatus(pset.db);
1071
1072                 /* We always rollback on an error */
1073                 if (transaction_status == PQTRANS_INERROR)
1074                         results = PQexec(pset.db, "ROLLBACK TO pg_psql_temporary_savepoint");
1075                 /* If they are no longer in a transaction, then do nothing */
1076                 else if (transaction_status != PQTRANS_INTRANS)
1077                         results = NULL;
1078                 else
1079                 {
1080                         /*
1081                          * Do nothing if they are messing with savepoints themselves: If
1082                          * the user did RELEASE or ROLLBACK, our savepoint is gone. If
1083                          * they issued a SAVEPOINT, releasing ours would remove theirs.
1084                          */
1085                         if (strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
1086                                 strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
1087                                 strcmp(PQcmdStatus(results), "ROLLBACK") == 0)
1088                                 results = NULL;
1089                         else
1090                                 results = PQexec(pset.db, "RELEASE pg_psql_temporary_savepoint");
1091                 }
1092                 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1093                 {
1094                         psql_error("%s", PQerrorMessage(pset.db));
1095                         PQclear(results);
1096                         ResetCancelConn();
1097                         return false;
1098                 }
1099                 PQclear(results);
1100         }
1101
1102         /* Possible microtiming output */
1103         if (OK && pset.timing)
1104                 printf(_("Time: %.3f ms\n"), DIFF_MSEC(&after, &before));
1105
1106         /* check for events that may occur during query execution */
1107
1108         if (pset.encoding != PQclientEncoding(pset.db) &&
1109                 PQclientEncoding(pset.db) >= 0)
1110         {
1111                 /* track effects of SET CLIENT_ENCODING */
1112                 pset.encoding = PQclientEncoding(pset.db);
1113                 pset.popt.topt.encoding = pset.encoding;
1114                 SetVariable(pset.vars, "ENCODING",
1115                                         pg_encoding_to_char(pset.encoding));
1116         }
1117
1118         PrintNotifications();
1119
1120         return OK;
1121 }
1122
1123
1124 /*
1125  * Advance the given char pointer over white space and SQL comments.
1126  */
1127 static const char *
1128 skip_white_space(const char *query)
1129 {
1130         int                     cnestlevel = 0; /* slash-star comment nest level */
1131
1132         while (*query)
1133         {
1134                 int                     mblen = PQmblen(query, pset.encoding);
1135
1136                 /*
1137                  * Note: we assume the encoding is a superset of ASCII, so that for
1138                  * example "query[0] == '/'" is meaningful.  However, we do NOT assume
1139                  * that the second and subsequent bytes of a multibyte character
1140                  * couldn't look like ASCII characters; so it is critical to advance
1141                  * by mblen, not 1, whenever we haven't exactly identified the
1142                  * character we are skipping over.
1143                  */
1144                 if (isspace((unsigned char) *query))
1145                         query += mblen;
1146                 else if (query[0] == '/' && query[1] == '*')
1147                 {
1148                         cnestlevel++;
1149                         query += 2;
1150                 }
1151                 else if (cnestlevel > 0 && query[0] == '*' && query[1] == '/')
1152                 {
1153                         cnestlevel--;
1154                         query += 2;
1155                 }
1156                 else if (cnestlevel == 0 && query[0] == '-' && query[1] == '-')
1157                 {
1158                         query += 2;
1159
1160                         /*
1161                          * We have to skip to end of line since any slash-star inside the
1162                          * -- comment does NOT start a slash-star comment.
1163                          */
1164                         while (*query)
1165                         {
1166                                 if (*query == '\n')
1167                                 {
1168                                         query++;
1169                                         break;
1170                                 }
1171                                 query += PQmblen(query, pset.encoding);
1172                         }
1173                 }
1174                 else if (cnestlevel > 0)
1175                         query += mblen;
1176                 else
1177                         break;                          /* found first token */
1178         }
1179
1180         return query;
1181 }
1182
1183
1184 /*
1185  * Check whether a command is one of those for which we should NOT start
1186  * a new transaction block (ie, send a preceding BEGIN).
1187  *
1188  * These include the transaction control statements themselves, plus
1189  * certain statements that the backend disallows inside transaction blocks.
1190  */
1191 static bool
1192 command_no_begin(const char *query)
1193 {
1194         int                     wordlen;
1195
1196         /*
1197          * First we must advance over any whitespace and comments.
1198          */
1199         query = skip_white_space(query);
1200
1201         /*
1202          * Check word length (since "beginx" is not "begin").
1203          */
1204         wordlen = 0;
1205         while (isalpha((unsigned char) query[wordlen]))
1206                 wordlen += PQmblen(&query[wordlen], pset.encoding);
1207
1208         /*
1209          * Transaction control commands.  These should include every keyword that
1210          * gives rise to a TransactionStmt in the backend grammar, except for the
1211          * savepoint-related commands.
1212          *
1213          * (We assume that START must be START TRANSACTION, since there is presently
1214          * no other "START foo" command.)
1215          */
1216         if (wordlen == 5 && pg_strncasecmp(query, "abort", 5) == 0)
1217                 return true;
1218         if (wordlen == 5 && pg_strncasecmp(query, "begin", 5) == 0)
1219                 return true;
1220         if (wordlen == 5 && pg_strncasecmp(query, "start", 5) == 0)
1221                 return true;
1222         if (wordlen == 6 && pg_strncasecmp(query, "commit", 6) == 0)
1223                 return true;
1224         if (wordlen == 3 && pg_strncasecmp(query, "end", 3) == 0)
1225                 return true;
1226         if (wordlen == 8 && pg_strncasecmp(query, "rollback", 8) == 0)
1227                 return true;
1228         if (wordlen == 7 && pg_strncasecmp(query, "prepare", 7) == 0)
1229         {
1230                 /* PREPARE TRANSACTION is a TC command, PREPARE foo is not */
1231                 query += wordlen;
1232
1233                 query = skip_white_space(query);
1234
1235                 wordlen = 0;
1236                 while (isalpha((unsigned char) query[wordlen]))
1237                         wordlen += PQmblen(&query[wordlen], pset.encoding);
1238
1239                 if (wordlen == 11 && pg_strncasecmp(query, "transaction", 11) == 0)
1240                         return true;
1241                 return false;
1242         }
1243
1244         /*
1245          * Commands not allowed within transactions.  The statements checked for
1246          * here should be exactly those that call PreventTransactionChain() in the
1247          * backend.
1248          *
1249          * Note: we are a bit sloppy about CLUSTER, which is transactional in some
1250          * variants but not others.
1251          */
1252         if (wordlen == 6 && pg_strncasecmp(query, "vacuum", 6) == 0)
1253                 return true;
1254         if (wordlen == 7 && pg_strncasecmp(query, "cluster", 7) == 0)
1255                 return true;
1256
1257         /*
1258          * Note: these tests will match CREATE SYSTEM, DROP SYSTEM, and REINDEX
1259          * TABLESPACE, which aren't really valid commands so we don't care much.
1260          * The other six possible matches are correct.
1261          */
1262         if ((wordlen == 6 && pg_strncasecmp(query, "create", 6) == 0) ||
1263                 (wordlen == 4 && pg_strncasecmp(query, "drop", 4) == 0) ||
1264                 (wordlen == 7 && pg_strncasecmp(query, "reindex", 7) == 0))
1265         {
1266                 query += wordlen;
1267
1268                 query = skip_white_space(query);
1269
1270                 wordlen = 0;
1271                 while (isalpha((unsigned char) query[wordlen]))
1272                         wordlen += PQmblen(&query[wordlen], pset.encoding);
1273
1274                 if (wordlen == 8 && pg_strncasecmp(query, "database", 8) == 0)
1275                         return true;
1276                 if (wordlen == 6 && pg_strncasecmp(query, "system", 6) == 0)
1277                         return true;
1278                 if (wordlen == 10 && pg_strncasecmp(query, "tablespace", 10) == 0)
1279                         return true;
1280         }
1281
1282         return false;
1283 }
1284
1285
1286 /*
1287  * Test if the current user is a database superuser.
1288  *
1289  * Note: this will correctly detect superuserness only with a protocol-3.0
1290  * or newer backend; otherwise it will always say "false".
1291  */
1292 bool
1293 is_superuser(void)
1294 {
1295         const char *val;
1296
1297         if (!pset.db)
1298                 return false;
1299
1300         val = PQparameterStatus(pset.db, "is_superuser");
1301
1302         if (val && strcmp(val, "on") == 0)
1303                 return true;
1304
1305         return false;
1306 }
1307
1308
1309 /*
1310  * Return the session user of the current connection.
1311  *
1312  * Note: this will correctly detect the session user only with a
1313  * protocol-3.0 or newer backend; otherwise it will return the
1314  * connection user.
1315  */
1316 const char *
1317 session_username(void)
1318 {
1319         const char *val;
1320
1321         if (!pset.db)
1322                 return NULL;
1323
1324         val = PQparameterStatus(pset.db, "session_authorization");
1325         if (val)
1326                 return val;
1327         else
1328                 return PQuser(pset.db);
1329 }
1330
1331
1332 /* expand_tilde
1333  *
1334  * substitute '~' with HOME or '~username' with username's home dir
1335  *
1336  */
1337 char *
1338 expand_tilde(char **filename)
1339 {
1340         if (!filename || !(*filename))
1341                 return NULL;
1342
1343         /*
1344          * WIN32 doesn't use tilde expansion for file names. Also, it uses tilde
1345          * for short versions of long file names, though the tilde is usually
1346          * toward the end, not at the beginning.
1347          */
1348 #ifndef WIN32
1349
1350         /* try tilde expansion */
1351         if (**filename == '~')
1352         {
1353                 char       *fn;
1354                 char            oldp,
1355                                    *p;
1356                 struct passwd *pw;
1357                 char            home[MAXPGPATH];
1358
1359                 fn = *filename;
1360                 *home = '\0';
1361
1362                 p = fn + 1;
1363                 while (*p != '/' && *p != '\0')
1364                         p++;
1365
1366                 oldp = *p;
1367                 *p = '\0';
1368
1369                 if (*(fn + 1) == '\0')
1370                         get_home_path(home);    /* ~ or ~/ only */
1371                 else if ((pw = getpwnam(fn + 1)) != NULL)
1372                         StrNCpy(home, pw->pw_dir, MAXPGPATH);           /* ~user */
1373
1374                 *p = oldp;
1375                 if (strlen(home) != 0)
1376                 {
1377                         char       *newfn;
1378
1379                         newfn = pg_malloc(strlen(home) + strlen(p) + 1);
1380                         strcpy(newfn, home);
1381                         strcat(newfn, p);
1382
1383                         free(fn);
1384                         *filename = newfn;
1385                 }
1386         }
1387 #endif
1388
1389         return *filename;
1390 }