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

   1 /*
   2  * psql - the PostgreSQL interactive terminal
   3  *
   4  * Copyright (c) 2000-2019, PostgreSQL Global Development Group
   5  *
   6  * src/bin/psql/common.c
   7  */
   8 #include "postgres_fe.h"
   9
  10 #include <ctype.h>
  11 #include <limits.h>
  12 #include <math.h>
  13 #include <signal.h>
  14 #ifndef WIN32
  15 #include <unistd.h>                             /* for write() */
  16 #else
  17 #include <io.h>                                 /* for _write() */
  18 #include <win32.h>
  19 #endif
  20
  21 #include "common/logging.h"
  22 #include "fe_utils/mbprint.h"
  23 #include "fe_utils/string_utils.h"
  24 #include "portability/instr_time.h"
  25
  26 #include "command.h"
  27 #include "common.h"
  28 #include "copy.h"
  29 #include "crosstabview.h"
  30 #include "settings.h"
  31
  32
  33 static bool DescribeQuery(const char *query, double *elapsed_msec);
  34 static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
  35 static bool command_no_begin(const char *query);
  36 static bool is_select_command(const char *query);
  37
  38
  39 /*
  40  * openQueryOutputFile --- attempt to open a query output file
  41  *
  42  * fname == NULL selects stdout, else an initial '|' selects a pipe,
  43  * else plain file.
  44  *
  45  * Returns output file pointer into *fout, and is-a-pipe flag into *is_pipe.
  46  * Caller is responsible for adjusting SIGPIPE state if it's a pipe.
  47  *
  48  * On error, reports suitable error message and returns false.
  49  */
  50 bool
  51 openQueryOutputFile(const char *fname, FILE **fout, bool *is_pipe)
  52 {
  53         if (!fname || fname[0] == '\0')
  54         {
  55                 *fout = stdout;
  56                 *is_pipe = false;
  57         }
  58         else if (*fname == '|')
  59         {
  60                 *fout = popen(fname + 1, "w");
  61                 *is_pipe = true;
  62         }
  63         else
  64         {
  65                 *fout = fopen(fname, "w");
  66                 *is_pipe = false;
  67         }
  68
  69         if (*fout == NULL)
  70         {
  71                 pg_log_error("%s: %m", fname);
  72                 return false;
  73         }
  74
  75         return true;
  76 }
  77
  78 /*
  79  * setQFout
  80  * -- handler for -o command line option and \o command
  81  *
  82  * On success, updates pset with the new output file and returns true.
  83  * On failure, returns false without changing pset state.
  84  */
  85 bool
  86 setQFout(const char *fname)
  87 {
  88         FILE       *fout;
  89         bool            is_pipe;
  90
  91         /* First make sure we can open the new output file/pipe */
  92         if (!openQueryOutputFile(fname, &fout, &is_pipe))
  93                 return false;
  94
  95         /* Close old file/pipe */
  96         if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
  97         {
  98                 if (pset.queryFoutPipe)
  99                         pclose(pset.queryFout);
 100                 else
 101                         fclose(pset.queryFout);
 102         }
 103
 104         pset.queryFout = fout;
 105         pset.queryFoutPipe = is_pipe;
 106
 107         /* Adjust SIGPIPE handling appropriately: ignore signal if is_pipe */
 108         set_sigpipe_trap_state(is_pipe);
 109         restore_sigpipe_trap();
 110
 111         return true;
 112 }
 113
 114
 115 /*
 116  * Variable-fetching callback for flex lexer
 117  *
 118  * If the specified variable exists, return its value as a string (malloc'd
 119  * and expected to be freed by the caller); else return NULL.
 120  *
 121  * If "quote" isn't PQUOTE_PLAIN, then return the value suitably quoted and
 122  * escaped for the specified quoting requirement.  (Failure in escaping
 123  * should lead to printing an error and returning NULL.)
 124  *
 125  * "passthrough" is the pointer previously given to psql_scan_set_passthrough.
 126  * In psql, passthrough points to a ConditionalStack, which we check to
 127  * determine whether variable expansion is allowed.
 128  */
 129 char *
 130 psql_get_variable(const char *varname, PsqlScanQuoteType quote,
 131                                   void *passthrough)
 132 {
 133         char       *result = NULL;
 134         const char *value;
 135
 136         /* In an inactive \if branch, suppress all variable substitutions */
 137         if (passthrough && !conditional_active((ConditionalStack) passthrough))
 138                 return NULL;
 139
 140         value = GetVariable(pset.vars, varname);
 141         if (!value)
 142                 return NULL;
 143
 144         switch (quote)
 145         {
 146                 case PQUOTE_PLAIN:
 147                         result = pg_strdup(value);
 148                         break;
 149                 case PQUOTE_SQL_LITERAL:
 150                 case PQUOTE_SQL_IDENT:
 151                         {
 152                                 /*
 153                                  * For these cases, we use libpq's quoting functions, which
 154                                  * assume the string is in the connection's client encoding.
 155                                  */
 156                                 char       *escaped_value;
 157
 158                                 if (!pset.db)
 159                                 {
 160                                         pg_log_error("cannot escape without active connection");
 161                                         return NULL;
 162                                 }
 163
 164                                 if (quote == PQUOTE_SQL_LITERAL)
 165                                         escaped_value =
 166                                                 PQescapeLiteral(pset.db, value, strlen(value));
 167                                 else
 168                                         escaped_value =
 169                                                 PQescapeIdentifier(pset.db, value, strlen(value));
 170
 171                                 if (escaped_value == NULL)
 172                                 {
 173                                         const char *error = PQerrorMessage(pset.db);
 174
 175                                         pg_log_info("%s", error);
 176                                         return NULL;
 177                                 }
 178
 179                                 /*
 180                                  * Rather than complicate the lexer's API with a notion of
 181                                  * which free() routine to use, just pay the price of an extra
 182                                  * strdup().
 183                                  */
 184                                 result = pg_strdup(escaped_value);
 185                                 PQfreemem(escaped_value);
 186                                 break;
 187                         }
 188                 case PQUOTE_SHELL_ARG:
 189                         {
 190                                 /*
 191                                  * For this we use appendShellStringNoError, which is
 192                                  * encoding-agnostic, which is fine since the shell probably
 193                                  * is too.  In any case, the only special character is "'",
 194                                  * which is not known to appear in valid multibyte characters.
 195                                  */
 196                                 PQExpBufferData buf;
 197
 198                                 initPQExpBuffer(&buf);
 199                                 if (!appendShellStringNoError(&buf, value))
 200                                 {
 201                                         pg_log_error("shell command argument contains a newline or carriage return: \"%s\"",
 202                                                                  value);
 203                                         free(buf.data);
 204                                         return NULL;
 205                                 }
 206                                 result = buf.data;
 207                                 break;
 208                         }
 209
 210                         /* No default: we want a compiler warning for missing cases */
 211         }
 212
 213         return result;
 214 }
 215
 216
 217 /*
 218  * for backend Notice messages (INFO, WARNING, etc)
 219  */
 220 void
 221 NoticeProcessor(void *arg, const char *message)
 222 {
 223         (void) arg;                                     /* not used */
 224         pg_log_info("%s", message);
 225 }
 226
 227
 228
 229 /*
 230  * Code to support query cancellation
 231  *
 232  * Before we start a query, we enable the SIGINT signal catcher to send a
 233  * cancel request to the backend. Note that sending the cancel directly from
 234  * the signal handler is safe because PQcancel() is written to make it
 235  * so. We use write() to report to stderr because it's better to use simple
 236  * facilities in a signal handler.
 237  *
 238  * On win32, the signal canceling happens on a separate thread, because
 239  * that's how SetConsoleCtrlHandler works. The PQcancel function is safe
 240  * for this (unlike PQrequestCancel). However, a CRITICAL_SECTION is required
 241  * to protect the PGcancel structure against being changed while the signal
 242  * thread is using it.
 243  *
 244  * SIGINT is supposed to abort all long-running psql operations, not only
 245  * database queries.  In most places, this is accomplished by checking
 246  * cancel_pressed during long-running loops.  However, that won't work when
 247  * blocked on user input (in readline() or fgets()).  In those places, we
 248  * set sigint_interrupt_enabled true while blocked, instructing the signal
 249  * catcher to longjmp through sigint_interrupt_jmp.  We assume readline and
 250  * fgets are coded to handle possible interruption.  (XXX currently this does
 251  * not work on win32, so control-C is less useful there)
 252  */
 253 volatile bool sigint_interrupt_enabled = false;
 254
 255 sigjmp_buf      sigint_interrupt_jmp;
 256
 257 static PGcancel *volatile cancelConn = NULL;
 258
 259 #ifdef WIN32
 260 static CRITICAL_SECTION cancelConnLock;
 261 #endif
 262
 263 /*
 264  * Write a simple string to stderr --- must be safe in a signal handler.
 265  * We ignore the write() result since there's not much we could do about it.
 266  * Certain compilers make that harder than it ought to be.
 267  */
 268 #define write_stderr(str) \
 269         do { \
 270                 const char *str_ = (str); \
 271                 int             rc_; \
 272                 rc_ = write(fileno(stderr), str_, strlen(str_)); \
 273                 (void) rc_; \
 274         } while (0)
 275
 276
 277 #ifndef WIN32
 278
 279 static void
 280 handle_sigint(SIGNAL_ARGS)
 281 {
 282         int                     save_errno = errno;
 283         char            errbuf[256];
 284
 285         /* if we are waiting for input, longjmp out of it */
 286         if (sigint_interrupt_enabled)
 287         {
 288                 sigint_interrupt_enabled = false;
 289                 siglongjmp(sigint_interrupt_jmp, 1);
 290         }
 291
 292         /* else, set cancel flag to stop any long-running loops */
 293         cancel_pressed = true;
 294
 295         /* and send QueryCancel if we are processing a database query */
 296         if (cancelConn != NULL)
 297         {
 298                 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 299                         write_stderr("Cancel request sent\n");
 300                 else
 301                 {
 302                         write_stderr("Could not send cancel request: ");
 303                         write_stderr(errbuf);
 304                 }
 305         }
 306
 307         errno = save_errno;                     /* just in case the write changed it */
 308 }
 309
 310 void
 311 setup_cancel_handler(void)
 312 {
 313         pqsignal(SIGINT, handle_sigint);
 314 }
 315 #else                                                   /* WIN32 */
 316
 317 static BOOL WINAPI
 318 consoleHandler(DWORD dwCtrlType)
 319 {
 320         char            errbuf[256];
 321
 322         if (dwCtrlType == CTRL_C_EVENT ||
 323                 dwCtrlType == CTRL_BREAK_EVENT)
 324         {
 325                 /*
 326                  * Can't longjmp here, because we are in wrong thread :-(
 327                  */
 328
 329                 /* set cancel flag to stop any long-running loops */
 330                 cancel_pressed = true;
 331
 332                 /* and send QueryCancel if we are processing a database query */
 333                 EnterCriticalSection(&cancelConnLock);
 334                 if (cancelConn != NULL)
 335                 {
 336                         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 337                                 write_stderr("Cancel request sent\n");
 338                         else
 339                         {
 340                                 write_stderr("Could not send cancel request: ");
 341                                 write_stderr(errbuf);
 342                         }
 343                 }
 344                 LeaveCriticalSection(&cancelConnLock);
 345
 346                 return TRUE;
 347         }
 348         else
 349                 /* Return FALSE for any signals not being handled */
 350                 return FALSE;
 351 }
 352
 353 void
 354 setup_cancel_handler(void)
 355 {
 356         InitializeCriticalSection(&cancelConnLock);
 357
 358         SetConsoleCtrlHandler(consoleHandler, TRUE);
 359 }
 360 #endif                                                  /* WIN32 */
 361
 362
 363 /* ConnectionUp
 364  *
 365  * Returns whether our backend connection is still there.
 366  */
 367 static bool
 368 ConnectionUp(void)
 369 {
 370         return PQstatus(pset.db) != CONNECTION_BAD;
 371 }
 372
 373
 374
 375 /* CheckConnection
 376  *
 377  * Verify that we still have a good connection to the backend, and if not,
 378  * see if it can be restored.
 379  *
 380  * Returns true if either the connection was still there, or it could be
 381  * restored successfully; false otherwise.  If, however, there was no
 382  * connection and the session is non-interactive, this will exit the program
 383  * with a code of EXIT_BADCONN.
 384  */
 385 static bool
 386 CheckConnection(void)
 387 {
 388         bool            OK;
 389
 390         OK = ConnectionUp();
 391         if (!OK)
 392         {
 393                 if (!pset.cur_cmd_interactive)
 394                 {
 395                         pg_log_fatal("connection to server was lost");
 396                         exit(EXIT_BADCONN);
 397                 }
 398
 399                 fprintf(stderr, _("The connection to the server was lost. Attempting reset: "));
 400                 PQreset(pset.db);
 401                 OK = ConnectionUp();
 402                 if (!OK)
 403                 {
 404                         fprintf(stderr, _("Failed.\n"));
 405                         PQfinish(pset.db);
 406                         pset.db = NULL;
 407                         ResetCancelConn();
 408                         UnsyncVariables();
 409                 }
 410                 else
 411                         fprintf(stderr, _("Succeeded.\n"));
 412         }
 413
 414         return OK;
 415 }
 416
 417
 418
 419 /*
 420  * SetCancelConn
 421  *
 422  * Set cancelConn to point to the current database connection.
 423  */
 424 void
 425 SetCancelConn(void)
 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         /* be sure handle_sigint doesn't use pointer while freeing */
 436         cancelConn = NULL;
 437
 438         if (oldCancelConn != NULL)
 439                 PQfreeCancel(oldCancelConn);
 440
 441         cancelConn = PQgetCancel(pset.db);
 442
 443 #ifdef WIN32
 444         LeaveCriticalSection(&cancelConnLock);
 445 #endif
 446 }
 447
 448
 449 /*
 450  * ResetCancelConn
 451  *
 452  * Free the current cancel connection, if any, and set to NULL.
 453  */
 454 void
 455 ResetCancelConn(void)
 456 {
 457         PGcancel   *oldCancelConn;
 458
 459 #ifdef WIN32
 460         EnterCriticalSection(&cancelConnLock);
 461 #endif
 462
 463         oldCancelConn = cancelConn;
 464         /* be sure handle_sigint doesn't use pointer while freeing */
 465         cancelConn = NULL;
 466
 467         if (oldCancelConn != NULL)
 468                 PQfreeCancel(oldCancelConn);
 469
 470 #ifdef WIN32
 471         LeaveCriticalSection(&cancelConnLock);
 472 #endif
 473 }
 474
 475
 476 /*
 477  * AcceptResult
 478  *
 479  * Checks whether a result is valid, giving an error message if necessary;
 480  * and ensures that the connection to the backend is still up.
 481  *
 482  * Returns true for valid result, false for error state.
 483  */
 484 static bool
 485 AcceptResult(const PGresult *result)
 486 {
 487         bool            OK;
 488
 489         if (!result)
 490                 OK = false;
 491         else
 492                 switch (PQresultStatus(result))
 493                 {
 494                         case PGRES_COMMAND_OK:
 495                         case PGRES_TUPLES_OK:
 496                         case PGRES_EMPTY_QUERY:
 497                         case PGRES_COPY_IN:
 498                         case PGRES_COPY_OUT:
 499                                 /* Fine, do nothing */
 500                                 OK = true;
 501                                 break;
 502
 503                         case PGRES_BAD_RESPONSE:
 504                         case PGRES_NONFATAL_ERROR:
 505                         case PGRES_FATAL_ERROR:
 506                                 OK = false;
 507                                 break;
 508
 509                         default:
 510                                 OK = false;
 511                                 pg_log_error("unexpected PQresultStatus: %d",
 512                                                          PQresultStatus(result));
 513                                 break;
 514                 }
 515
 516         if (!OK)
 517         {
 518                 const char *error = PQerrorMessage(pset.db);
 519
 520                 if (strlen(error))
 521                         pg_log_info("%s", error);
 522
 523                 CheckConnection();
 524         }
 525
 526         return OK;
 527 }
 528
 529
 530 /*
 531  * Set special variables from a query result
 532  * - ERROR: true/false, whether an error occurred on this query
 533  * - SQLSTATE: code of error, or "00000" if no error, or "" if unknown
 534  * - ROW_COUNT: how many rows were returned or affected, or "0"
 535  * - LAST_ERROR_SQLSTATE: same for last error
 536  * - LAST_ERROR_MESSAGE: message of last error
 537  *
 538  * Note: current policy is to apply this only to the results of queries
 539  * entered by the user, not queries generated by slash commands.
 540  */
 541 static void
 542 SetResultVariables(PGresult *results, bool success)
 543 {
 544         if (success)
 545         {
 546                 const char *ntuples = PQcmdTuples(results);
 547
 548                 SetVariable(pset.vars, "ERROR", "false");
 549                 SetVariable(pset.vars, "SQLSTATE", "00000");
 550                 SetVariable(pset.vars, "ROW_COUNT", *ntuples ? ntuples : "0");
 551         }
 552         else
 553         {
 554                 const char *code = PQresultErrorField(results, PG_DIAG_SQLSTATE);
 555                 const char *mesg = PQresultErrorField(results, PG_DIAG_MESSAGE_PRIMARY);
 556
 557                 SetVariable(pset.vars, "ERROR", "true");
 558
 559                 /*
 560                  * If there is no SQLSTATE code, use an empty string.  This can happen
 561                  * for libpq-detected errors (e.g., lost connection, ENOMEM).
 562                  */
 563                 if (code == NULL)
 564                         code = "";
 565                 SetVariable(pset.vars, "SQLSTATE", code);
 566                 SetVariable(pset.vars, "ROW_COUNT", "0");
 567                 SetVariable(pset.vars, "LAST_ERROR_SQLSTATE", code);
 568                 SetVariable(pset.vars, "LAST_ERROR_MESSAGE", mesg ? mesg : "");
 569         }
 570 }
 571
 572
 573 /*
 574  * ClearOrSaveResult
 575  *
 576  * If the result represents an error, remember it for possible display by
 577  * \errverbose.  Otherwise, just PQclear() it.
 578  *
 579  * Note: current policy is to apply this to the results of all queries,
 580  * including "back door" queries, for debugging's sake.  It's OK to use
 581  * PQclear() directly on results known to not be error results, however.
 582  */
 583 static void
 584 ClearOrSaveResult(PGresult *result)
 585 {
 586         if (result)
 587         {
 588                 switch (PQresultStatus(result))
 589                 {
 590                         case PGRES_NONFATAL_ERROR:
 591                         case PGRES_FATAL_ERROR:
 592                                 if (pset.last_error_result)
 593                                         PQclear(pset.last_error_result);
 594                                 pset.last_error_result = result;
 595                                 break;
 596
 597                         default:
 598                                 PQclear(result);
 599                                 break;
 600                 }
 601         }
 602 }
 603
 604
 605 /*
 606  * Print microtiming output.  Always print raw milliseconds; if the interval
 607  * is >= 1 second, also break it down into days/hours/minutes/seconds.
 608  */
 609 static void
 610 PrintTiming(double elapsed_msec)
 611 {
 612         double          seconds;
 613         double          minutes;
 614         double          hours;
 615         double          days;
 616
 617         if (elapsed_msec < 1000.0)
 618         {
 619                 /* This is the traditional (pre-v10) output format */
 620                 printf(_("Time: %.3f ms\n"), elapsed_msec);
 621                 return;
 622         }
 623
 624         /*
 625          * Note: we could print just seconds, in a format like %06.3f, when the
 626          * total is less than 1min.  But that's hard to interpret unless we tack
 627          * on "s" or otherwise annotate it.  Forcing the display to include
 628          * minutes seems like a better solution.
 629          */
 630         seconds = elapsed_msec / 1000.0;
 631         minutes = floor(seconds / 60.0);
 632         seconds -= 60.0 * minutes;
 633         if (minutes < 60.0)
 634         {
 635                 printf(_("Time: %.3f ms (%02d:%06.3f)\n"),
 636                            elapsed_msec, (int) minutes, seconds);
 637                 return;
 638         }
 639
 640         hours = floor(minutes / 60.0);
 641         minutes -= 60.0 * hours;
 642         if (hours < 24.0)
 643         {
 644                 printf(_("Time: %.3f ms (%02d:%02d:%06.3f)\n"),
 645                            elapsed_msec, (int) hours, (int) minutes, seconds);
 646                 return;
 647         }
 648
 649         days = floor(hours / 24.0);
 650         hours -= 24.0 * days;
 651         printf(_("Time: %.3f ms (%.0f d %02d:%02d:%06.3f)\n"),
 652                    elapsed_msec, days, (int) hours, (int) minutes, seconds);
 653 }
 654
 655
 656 /*
 657  * PSQLexec
 658  *
 659  * This is the way to send "backdoor" queries (those not directly entered
 660  * by the user). It is subject to -E but not -e.
 661  *
 662  * Caller is responsible for handling the ensuing processing if a COPY
 663  * command is sent.
 664  *
 665  * Note: we don't bother to check PQclientEncoding; it is assumed that no
 666  * caller uses this path to issue "SET CLIENT_ENCODING".
 667  */
 668 PGresult *
 669 PSQLexec(const char *query)
 670 {
 671         PGresult   *res;
 672
 673         if (!pset.db)
 674         {
 675                 pg_log_error("You are currently not connected to a database.");
 676                 return NULL;
 677         }
 678
 679         if (pset.echo_hidden != PSQL_ECHO_HIDDEN_OFF)
 680         {
 681                 printf(_("********* QUERY **********\n"
 682                                  "%s\n"
 683                                  "**************************\n\n"), query);
 684                 fflush(stdout);
 685                 if (pset.logfile)
 686                 {
 687                         fprintf(pset.logfile,
 688                                         _("********* QUERY **********\n"
 689                                           "%s\n"
 690                                           "**************************\n\n"), query);
 691                         fflush(pset.logfile);
 692                 }
 693
 694                 if (pset.echo_hidden == PSQL_ECHO_HIDDEN_NOEXEC)
 695                         return NULL;
 696         }
 697
 698         SetCancelConn();
 699
 700         res = PQexec(pset.db, query);
 701
 702         ResetCancelConn();
 703
 704         if (!AcceptResult(res))
 705         {
 706                 ClearOrSaveResult(res);
 707                 res = NULL;
 708         }
 709
 710         return res;
 711 }
 712
 713
 714 /*
 715  * PSQLexecWatch
 716  *
 717  * This function is used for \watch command to send the query to
 718  * the server and print out the results.
 719  *
 720  * Returns 1 if the query executed successfully, 0 if it cannot be repeated,
 721  * e.g., because of the interrupt, -1 on error.
 722  */
 723 int
 724 PSQLexecWatch(const char *query, const printQueryOpt *opt)
 725 {
 726         PGresult   *res;
 727         double          elapsed_msec = 0;
 728         instr_time      before;
 729         instr_time      after;
 730
 731         if (!pset.db)
 732         {
 733                 pg_log_error("You are currently not connected to a database.");
 734                 return 0;
 735         }
 736
 737         SetCancelConn();
 738
 739         if (pset.timing)
 740                 INSTR_TIME_SET_CURRENT(before);
 741
 742         res = PQexec(pset.db, query);
 743
 744         ResetCancelConn();
 745
 746         if (!AcceptResult(res))
 747         {
 748                 ClearOrSaveResult(res);
 749                 return 0;
 750         }
 751
 752         if (pset.timing)
 753         {
 754                 INSTR_TIME_SET_CURRENT(after);
 755                 INSTR_TIME_SUBTRACT(after, before);
 756                 elapsed_msec = INSTR_TIME_GET_MILLISEC(after);
 757         }
 758
 759         /*
 760          * If SIGINT is sent while the query is processing, the interrupt will be
 761          * consumed.  The user's intention, though, is to cancel the entire watch
 762          * process, so detect a sent cancellation request and exit in this case.
 763          */
 764         if (cancel_pressed)
 765         {
 766                 PQclear(res);
 767                 return 0;
 768         }
 769
 770         switch (PQresultStatus(res))
 771         {
 772                 case PGRES_TUPLES_OK:
 773                         printQuery(res, opt, pset.queryFout, false, pset.logfile);
 774                         break;
 775
 776                 case PGRES_COMMAND_OK:
 777                         fprintf(pset.queryFout, "%s\n%s\n\n", opt->title, PQcmdStatus(res));
 778                         break;
 779
 780                 case PGRES_EMPTY_QUERY:
 781                         pg_log_error("\\watch cannot be used with an empty query");
 782                         PQclear(res);
 783                         return -1;
 784
 785                 case PGRES_COPY_OUT:
 786                 case PGRES_COPY_IN:
 787                 case PGRES_COPY_BOTH:
 788                         pg_log_error("\\watch cannot be used with COPY");
 789                         PQclear(res);
 790                         return -1;
 791
 792                 default:
 793                         pg_log_error("unexpected result status for \\watch");
 794                         PQclear(res);
 795                         return -1;
 796         }
 797
 798         PQclear(res);
 799
 800         fflush(pset.queryFout);
 801
 802         /* Possible microtiming output */
 803         if (pset.timing)
 804                 PrintTiming(elapsed_msec);
 805
 806         return 1;
 807 }
 808
 809
 810 /*
 811  * PrintNotifications: check for asynchronous notifications, and print them out
 812  */
 813 static void
 814 PrintNotifications(void)
 815 {
 816         PGnotify   *notify;
 817
 818         PQconsumeInput(pset.db);
 819         while ((notify = PQnotifies(pset.db)) != NULL)
 820         {
 821                 /* for backward compatibility, only show payload if nonempty */
 822                 if (notify->extra[0])
 823                         fprintf(pset.queryFout, _("Asynchronous notification \"%s\" with payload \"%s\" received from server process with PID %d.\n"),
 824                                         notify->relname, notify->extra, notify->be_pid);
 825                 else
 826                         fprintf(pset.queryFout, _("Asynchronous notification \"%s\" received from server process with PID %d.\n"),
 827                                         notify->relname, notify->be_pid);
 828                 fflush(pset.queryFout);
 829                 PQfreemem(notify);
 830                 PQconsumeInput(pset.db);
 831         }
 832 }
 833
 834
 835 /*
 836  * PrintQueryTuples: assuming query result is OK, print its tuples
 837  *
 838  * Returns true if successful, false otherwise.
 839  */
 840 static bool
 841 PrintQueryTuples(const PGresult *results)
 842 {
 843         printQueryOpt my_popt = pset.popt;
 844
 845         /* one-shot expanded output requested via \gx */
 846         if (pset.g_expanded)
 847                 my_popt.topt.expanded = 1;
 848
 849         /* write output to \g argument, if any */
 850         if (pset.gfname)
 851         {
 852                 FILE       *fout;
 853                 bool            is_pipe;
 854
 855                 if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
 856                         return false;
 857                 if (is_pipe)
 858                         disable_sigpipe_trap();
 859
 860                 printQuery(results, &my_popt, fout, false, pset.logfile);
 861
 862                 if (is_pipe)
 863                 {
 864                         pclose(fout);
 865                         restore_sigpipe_trap();
 866                 }
 867                 else
 868                         fclose(fout);
 869         }
 870         else
 871                 printQuery(results, &my_popt, pset.queryFout, false, pset.logfile);
 872
 873         return true;
 874 }
 875
 876
 877 /*
 878  * StoreQueryTuple: assuming query result is OK, save data into variables
 879  *
 880  * Returns true if successful, false otherwise.
 881  */
 882 static bool
 883 StoreQueryTuple(const PGresult *result)
 884 {
 885         bool            success = true;
 886
 887         if (PQntuples(result) < 1)
 888         {
 889                 pg_log_error("no rows returned for \\gset");
 890                 success = false;
 891         }
 892         else if (PQntuples(result) > 1)
 893         {
 894                 pg_log_error("more than one row returned for \\gset");
 895                 success = false;
 896         }
 897         else
 898         {
 899                 int                     i;
 900
 901                 for (i = 0; i < PQnfields(result); i++)
 902                 {
 903                         char       *colname = PQfname(result, i);
 904                         char       *varname;
 905                         char       *value;
 906
 907                         /* concatenate prefix and column name */
 908                         varname = psprintf("%s%s", pset.gset_prefix, colname);
 909
 910                         if (!PQgetisnull(result, 0, i))
 911                                 value = PQgetvalue(result, 0, i);
 912                         else
 913                         {
 914                                 /* for NULL value, unset rather than set the variable */
 915                                 value = NULL;
 916                         }
 917
 918                         if (!SetVariable(pset.vars, varname, value))
 919                         {
 920                                 free(varname);
 921                                 success = false;
 922                                 break;
 923                         }
 924
 925                         free(varname);
 926                 }
 927         }
 928
 929         return success;
 930 }
 931
 932
 933 /*
 934  * ExecQueryTuples: assuming query result is OK, execute each query
 935  * result field as a SQL statement
 936  *
 937  * Returns true if successful, false otherwise.
 938  */
 939 static bool
 940 ExecQueryTuples(const PGresult *result)
 941 {
 942         bool            success = true;
 943         int                     nrows = PQntuples(result);
 944         int                     ncolumns = PQnfields(result);
 945         int                     r,
 946                                 c;
 947
 948         /*
 949          * We must turn off gexec_flag to avoid infinite recursion.  Note that
 950          * this allows ExecQueryUsingCursor to be applied to the individual query
 951          * results.  SendQuery prevents it from being applied when fetching the
 952          * queries-to-execute, because it can't handle recursion either.
 953          */
 954         pset.gexec_flag = false;
 955
 956         for (r = 0; r < nrows; r++)
 957         {
 958                 for (c = 0; c < ncolumns; c++)
 959                 {
 960                         if (!PQgetisnull(result, r, c))
 961                         {
 962                                 const char *query = PQgetvalue(result, r, c);
 963
 964                                 /* Abandon execution if cancel_pressed */
 965                                 if (cancel_pressed)
 966                                         goto loop_exit;
 967
 968                                 /*
 969                                  * ECHO_ALL mode should echo these queries, but SendQuery
 970                                  * assumes that MainLoop did that, so we have to do it here.
 971                                  */
 972                                 if (pset.echo == PSQL_ECHO_ALL && !pset.singlestep)
 973                                 {
 974                                         puts(query);
 975                                         fflush(stdout);
 976                                 }
 977
 978                                 if (!SendQuery(query))
 979                                 {
 980                                         /* Error - abandon execution if ON_ERROR_STOP */
 981                                         success = false;
 982                                         if (pset.on_error_stop)
 983                                                 goto loop_exit;
 984                                 }
 985                         }
 986                 }
 987         }
 988
 989 loop_exit:
 990
 991         /*
 992          * Restore state.  We know gexec_flag was on, else we'd not be here. (We
 993          * also know it'll get turned off at end of command, but that's not ours
 994          * to do here.)
 995          */
 996         pset.gexec_flag = true;
 997
 998         /* Return true if all queries were successful */
 999         return success;
1000 }
1001
1002
1003 /*
1004  * ProcessResult: utility function for use by SendQuery() only
1005  *
1006  * When our command string contained a COPY FROM STDIN or COPY TO STDOUT,
1007  * PQexec() has stopped at the PGresult associated with the first such
1008  * command.  In that event, we'll marshal data for the COPY and then cycle
1009  * through any subsequent PGresult objects.
1010  *
1011  * When the command string contained no such COPY command, this function
1012  * degenerates to an AcceptResult() call.
1013  *
1014  * Changes its argument to point to the last PGresult of the command string,
1015  * or NULL if that result was for a COPY TO STDOUT.  (Returning NULL prevents
1016  * the command status from being printed, which we want in that case so that
1017  * the status line doesn't get taken as part of the COPY data.)
1018  *
1019  * Returns true on complete success, false otherwise.  Possible failure modes
1020  * include purely client-side problems; check the transaction status for the
1021  * server-side opinion.
1022  */
1023 static bool
1024 ProcessResult(PGresult **results)
1025 {
1026         bool            success = true;
1027         bool            first_cycle = true;
1028
1029         for (;;)
1030         {
1031                 ExecStatusType result_status;
1032                 bool            is_copy;
1033                 PGresult   *next_result;
1034
1035                 if (!AcceptResult(*results))
1036                 {
1037                         /*
1038                          * Failure at this point is always a server-side failure or a
1039                          * failure to submit the command string.  Either way, we're
1040                          * finished with this command string.
1041                          */
1042                         success = false;
1043                         break;
1044                 }
1045
1046                 result_status = PQresultStatus(*results);
1047                 switch (result_status)
1048                 {
1049                         case PGRES_EMPTY_QUERY:
1050                         case PGRES_COMMAND_OK:
1051                         case PGRES_TUPLES_OK:
1052                                 is_copy = false;
1053                                 break;
1054
1055                         case PGRES_COPY_OUT:
1056                         case PGRES_COPY_IN:
1057                                 is_copy = true;
1058                                 break;
1059
1060                         default:
1061                                 /* AcceptResult() should have caught anything else. */
1062                                 is_copy = false;
1063                                 pg_log_error("unexpected PQresultStatus: %d", result_status);
1064                                 break;
1065                 }
1066
1067                 if (is_copy)
1068                 {
1069                         /*
1070                          * Marshal the COPY data.  Either subroutine will get the
1071                          * connection out of its COPY state, then call PQresultStatus()
1072                          * once and report any error.
1073                          *
1074                          * For COPY OUT, direct the output to pset.copyStream if it's set,
1075                          * otherwise to pset.gfname if it's set, otherwise to queryFout.
1076                          * For COPY IN, use pset.copyStream as data source if it's set,
1077                          * otherwise cur_cmd_source.
1078                          */
1079                         FILE       *copystream;
1080                         PGresult   *copy_result;
1081
1082                         SetCancelConn();
1083                         if (result_status == PGRES_COPY_OUT)
1084                         {
1085                                 bool            need_close = false;
1086                                 bool            is_pipe = false;
1087
1088                                 if (pset.copyStream)
1089                                 {
1090                                         /* invoked by \copy */
1091                                         copystream = pset.copyStream;
1092                                 }
1093                                 else if (pset.gfname)
1094                                 {
1095                                         /* invoked by \g */
1096                                         if (openQueryOutputFile(pset.gfname,
1097                                                                                         &copystream, &is_pipe))
1098                                         {
1099                                                 need_close = true;
1100                                                 if (is_pipe)
1101                                                         disable_sigpipe_trap();
1102                                         }
1103                                         else
1104                                                 copystream = NULL;      /* discard COPY data entirely */
1105                                 }
1106                                 else
1107                                 {
1108                                         /* fall back to the generic query output stream */
1109                                         copystream = pset.queryFout;
1110                                 }
1111
1112                                 success = handleCopyOut(pset.db,
1113                                                                                 copystream,
1114                                                                                 &copy_result)
1115                                         && success
1116                                         && (copystream != NULL);
1117
1118                                 /*
1119                                  * Suppress status printing if the report would go to the same
1120                                  * place as the COPY data just went.  Note this doesn't
1121                                  * prevent error reporting, since handleCopyOut did that.
1122                                  */
1123                                 if (copystream == pset.queryFout)
1124                                 {
1125                                         PQclear(copy_result);
1126                                         copy_result = NULL;
1127                                 }
1128
1129                                 if (need_close)
1130                                 {
1131                                         /* close \g argument file/pipe */
1132                                         if (is_pipe)
1133                                         {
1134                                                 pclose(copystream);
1135                                                 restore_sigpipe_trap();
1136                                         }
1137                                         else
1138                                         {
1139                                                 fclose(copystream);
1140                                         }
1141                                 }
1142                         }
1143                         else
1144                         {
1145                                 /* COPY IN */
1146                                 copystream = pset.copyStream ? pset.copyStream : pset.cur_cmd_source;
1147                                 success = handleCopyIn(pset.db,
1148                                                                            copystream,
1149                                                                            PQbinaryTuples(*results),
1150                                                                            &copy_result) && success;
1151                         }
1152                         ResetCancelConn();
1153
1154                         /*
1155                          * Replace the PGRES_COPY_OUT/IN result with COPY command's exit
1156                          * status, or with NULL if we want to suppress printing anything.
1157                          */
1158                         PQclear(*results);
1159                         *results = copy_result;
1160                 }
1161                 else if (first_cycle)
1162                 {
1163                         /* fast path: no COPY commands; PQexec visited all results */
1164                         break;
1165                 }
1166
1167                 /*
1168                  * Check PQgetResult() again.  In the typical case of a single-command
1169                  * string, it will return NULL.  Otherwise, we'll have other results
1170                  * to process that may include other COPYs.  We keep the last result.
1171                  */
1172                 next_result = PQgetResult(pset.db);
1173                 if (!next_result)
1174                         break;
1175
1176                 PQclear(*results);
1177                 *results = next_result;
1178                 first_cycle = false;
1179         }
1180
1181         SetResultVariables(*results, success);
1182
1183         /* may need this to recover from conn loss during COPY */
1184         if (!first_cycle && !CheckConnection())
1185                 return false;
1186
1187         return success;
1188 }
1189
1190
1191 /*
1192  * PrintQueryStatus: report command status as required
1193  *
1194  * Note: Utility function for use by PrintQueryResults() only.
1195  */
1196 static void
1197 PrintQueryStatus(PGresult *results)
1198 {
1199         char            buf[16];
1200
1201         if (!pset.quiet)
1202         {
1203                 if (pset.popt.topt.format == PRINT_HTML)
1204                 {
1205                         fputs("<p>", pset.queryFout);
1206                         html_escaped_print(PQcmdStatus(results), pset.queryFout);
1207                         fputs("</p>\n", pset.queryFout);
1208                 }
1209                 else
1210                         fprintf(pset.queryFout, "%s\n", PQcmdStatus(results));
1211         }
1212
1213         if (pset.logfile)
1214                 fprintf(pset.logfile, "%s\n", PQcmdStatus(results));
1215
1216         snprintf(buf, sizeof(buf), "%u", (unsigned int) PQoidValue(results));
1217         SetVariable(pset.vars, "LASTOID", buf);
1218 }
1219
1220
1221 /*
1222  * PrintQueryResults: print out (or store or execute) query results as required
1223  *
1224  * Note: Utility function for use by SendQuery() only.
1225  *
1226  * Returns true if the query executed successfully, false otherwise.
1227  */
1228 static bool
1229 PrintQueryResults(PGresult *results)
1230 {
1231         bool            success;
1232         const char *cmdstatus;
1233
1234         if (!results)
1235                 return false;
1236
1237         switch (PQresultStatus(results))
1238         {
1239                 case PGRES_TUPLES_OK:
1240                         /* store or execute or print the data ... */
1241                         if (pset.gset_prefix)
1242                                 success = StoreQueryTuple(results);
1243                         else if (pset.gexec_flag)
1244                                 success = ExecQueryTuples(results);
1245                         else if (pset.crosstab_flag)
1246                                 success = PrintResultsInCrosstab(results);
1247                         else
1248                                 success = PrintQueryTuples(results);
1249                         /* if it's INSERT/UPDATE/DELETE RETURNING, also print status */
1250                         cmdstatus = PQcmdStatus(results);
1251                         if (strncmp(cmdstatus, "INSERT", 6) == 0 ||
1252                                 strncmp(cmdstatus, "UPDATE", 6) == 0 ||
1253                                 strncmp(cmdstatus, "DELETE", 6) == 0)
1254                                 PrintQueryStatus(results);
1255                         break;
1256
1257                 case PGRES_COMMAND_OK:
1258                         PrintQueryStatus(results);
1259                         success = true;
1260                         break;
1261
1262                 case PGRES_EMPTY_QUERY:
1263                         success = true;
1264                         break;
1265
1266                 case PGRES_COPY_OUT:
1267                 case PGRES_COPY_IN:
1268                         /* nothing to do here */
1269                         success = true;
1270                         break;
1271
1272                 case PGRES_BAD_RESPONSE:
1273                 case PGRES_NONFATAL_ERROR:
1274                 case PGRES_FATAL_ERROR:
1275                         success = false;
1276                         break;
1277
1278                 default:
1279                         success = false;
1280                         pg_log_error("unexpected PQresultStatus: %d",
1281                                                  PQresultStatus(results));
1282                         break;
1283         }
1284
1285         fflush(pset.queryFout);
1286
1287         return success;
1288 }
1289
1290
1291 /*
1292  * SendQuery: send the query string to the backend
1293  * (and print out results)
1294  *
1295  * Note: This is the "front door" way to send a query. That is, use it to
1296  * send queries actually entered by the user. These queries will be subject to
1297  * single step mode.
1298  * To send "back door" queries (generated by slash commands, etc.) in a
1299  * controlled way, use PSQLexec().
1300  *
1301  * Returns true if the query executed successfully, false otherwise.
1302  */
1303 bool
1304 SendQuery(const char *query)
1305 {
1306         PGresult   *results;
1307         PGTransactionStatusType transaction_status;
1308         double          elapsed_msec = 0;
1309         bool            OK = false;
1310         int                     i;
1311         bool            on_error_rollback_savepoint = false;
1312         static bool on_error_rollback_warning = false;
1313
1314         if (!pset.db)
1315         {
1316                 pg_log_error("You are currently not connected to a database.");
1317                 goto sendquery_cleanup;
1318         }
1319
1320         if (pset.singlestep)
1321         {
1322                 char            buf[3];
1323
1324                 fflush(stderr);
1325                 printf(_("***(Single step mode: verify command)*******************************************\n"
1326                                  "%s\n"
1327                                  "***(press return to proceed or enter x and return to cancel)********************\n"),
1328                            query);
1329                 fflush(stdout);
1330                 if (fgets(buf, sizeof(buf), stdin) != NULL)
1331                         if (buf[0] == 'x')
1332                                 goto sendquery_cleanup;
1333                 if (cancel_pressed)
1334                         goto sendquery_cleanup;
1335         }
1336         else if (pset.echo == PSQL_ECHO_QUERIES)
1337         {
1338                 puts(query);
1339                 fflush(stdout);
1340         }
1341
1342         if (pset.logfile)
1343         {
1344                 fprintf(pset.logfile,
1345                                 _("********* QUERY **********\n"
1346                                   "%s\n"
1347                                   "**************************\n\n"), query);
1348                 fflush(pset.logfile);
1349         }
1350
1351         SetCancelConn();
1352
1353         transaction_status = PQtransactionStatus(pset.db);
1354
1355         if (transaction_status == PQTRANS_IDLE &&
1356                 !pset.autocommit &&
1357                 !command_no_begin(query))
1358         {
1359                 results = PQexec(pset.db, "BEGIN");
1360                 if (PQresultStatus(results) != PGRES_COMMAND_OK)
1361                 {
1362                         pg_log_info("%s", PQerrorMessage(pset.db));
1363                         ClearOrSaveResult(results);
1364                         ResetCancelConn();
1365                         goto sendquery_cleanup;
1366                 }
1367                 ClearOrSaveResult(results);
1368                 transaction_status = PQtransactionStatus(pset.db);
1369         }
1370
1371         if (transaction_status == PQTRANS_INTRANS &&
1372                 pset.on_error_rollback != PSQL_ERROR_ROLLBACK_OFF &&
1373                 (pset.cur_cmd_interactive ||
1374                  pset.on_error_rollback == PSQL_ERROR_ROLLBACK_ON))
1375         {
1376                 if (on_error_rollback_warning == false && pset.sversion < 80000)
1377                 {
1378                         char            sverbuf[32];
1379
1380                         pg_log_warning("The server (version %s) does not support savepoints for ON_ERROR_ROLLBACK.",
1381                                                    formatPGVersionNumber(pset.sversion, false,
1382                                                                                                  sverbuf, sizeof(sverbuf)));
1383                         on_error_rollback_warning = true;
1384                 }
1385                 else
1386                 {
1387                         results = PQexec(pset.db, "SAVEPOINT pg_psql_temporary_savepoint");
1388                         if (PQresultStatus(results) != PGRES_COMMAND_OK)
1389                         {
1390                                 pg_log_info("%s", PQerrorMessage(pset.db));
1391                                 ClearOrSaveResult(results);
1392                                 ResetCancelConn();
1393                                 goto sendquery_cleanup;
1394                         }
1395                         ClearOrSaveResult(results);
1396                         on_error_rollback_savepoint = true;
1397                 }
1398         }
1399
1400         if (pset.gdesc_flag)
1401         {
1402                 /* Describe query's result columns, without executing it */
1403                 OK = DescribeQuery(query, &elapsed_msec);
1404                 ResetCancelConn();
1405                 results = NULL;                 /* PQclear(NULL) does nothing */
1406         }
1407         else if (pset.fetch_count <= 0 || pset.gexec_flag ||
1408                          pset.crosstab_flag || !is_select_command(query))
1409         {
1410                 /* Default fetch-it-all-and-print mode */
1411                 instr_time      before,
1412                                         after;
1413
1414                 if (pset.timing)
1415                         INSTR_TIME_SET_CURRENT(before);
1416
1417                 results = PQexec(pset.db, query);
1418
1419                 /* these operations are included in the timing result: */
1420                 ResetCancelConn();
1421                 OK = ProcessResult(&results);
1422
1423                 if (pset.timing)
1424                 {
1425                         INSTR_TIME_SET_CURRENT(after);
1426                         INSTR_TIME_SUBTRACT(after, before);
1427                         elapsed_msec = INSTR_TIME_GET_MILLISEC(after);
1428                 }
1429
1430                 /* but printing results isn't: */
1431                 if (OK && results)
1432                         OK = PrintQueryResults(results);
1433         }
1434         else
1435         {
1436                 /* Fetch-in-segments mode */
1437                 OK = ExecQueryUsingCursor(query, &elapsed_msec);
1438                 ResetCancelConn();
1439                 results = NULL;                 /* PQclear(NULL) does nothing */
1440         }
1441
1442         if (!OK && pset.echo == PSQL_ECHO_ERRORS)
1443                 pg_log_info("STATEMENT:  %s", query);
1444
1445         /* If we made a temporary savepoint, possibly release/rollback */
1446         if (on_error_rollback_savepoint)
1447         {
1448                 const char *svptcmd = NULL;
1449
1450                 transaction_status = PQtransactionStatus(pset.db);
1451
1452                 switch (transaction_status)
1453                 {
1454                         case PQTRANS_INERROR:
1455                                 /* We always rollback on an error */
1456                                 svptcmd = "ROLLBACK TO pg_psql_temporary_savepoint";
1457                                 break;
1458
1459                         case PQTRANS_IDLE:
1460                                 /* If they are no longer in a transaction, then do nothing */
1461                                 break;
1462
1463                         case PQTRANS_INTRANS:
1464
1465                                 /*
1466                                  * Do nothing if they are messing with savepoints themselves:
1467                                  * If the user did RELEASE or ROLLBACK, our savepoint is gone.
1468                                  * If they issued a SAVEPOINT, releasing ours would remove
1469                                  * theirs.
1470                                  */
1471                                 if (results &&
1472                                         (strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
1473                                          strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
1474                                          strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
1475                                         svptcmd = NULL;
1476                                 else
1477                                         svptcmd = "RELEASE pg_psql_temporary_savepoint";
1478                                 break;
1479
1480                         case PQTRANS_ACTIVE:
1481                         case PQTRANS_UNKNOWN:
1482                         default:
1483                                 OK = false;
1484                                 /* PQTRANS_UNKNOWN is expected given a broken connection. */
1485                                 if (transaction_status != PQTRANS_UNKNOWN || ConnectionUp())
1486                                         pg_log_error("unexpected transaction status (%d)",
1487                                                                  transaction_status);
1488                                 break;
1489                 }
1490
1491                 if (svptcmd)
1492                 {
1493                         PGresult   *svptres;
1494
1495                         svptres = PQexec(pset.db, svptcmd);
1496                         if (PQresultStatus(svptres) != PGRES_COMMAND_OK)
1497                         {
1498                                 pg_log_info("%s", PQerrorMessage(pset.db));
1499                                 ClearOrSaveResult(svptres);
1500                                 OK = false;
1501
1502                                 PQclear(results);
1503                                 ResetCancelConn();
1504                                 goto sendquery_cleanup;
1505                         }
1506                         PQclear(svptres);
1507                 }
1508         }
1509
1510         ClearOrSaveResult(results);
1511
1512         /* Possible microtiming output */
1513         if (pset.timing)
1514                 PrintTiming(elapsed_msec);
1515
1516         /* check for events that may occur during query execution */
1517
1518         if (pset.encoding != PQclientEncoding(pset.db) &&
1519                 PQclientEncoding(pset.db) >= 0)
1520         {
1521                 /* track effects of SET CLIENT_ENCODING */
1522                 pset.encoding = PQclientEncoding(pset.db);
1523                 pset.popt.topt.encoding = pset.encoding;
1524                 SetVariable(pset.vars, "ENCODING",
1525                                         pg_encoding_to_char(pset.encoding));
1526         }
1527
1528         PrintNotifications();
1529
1530         /* perform cleanup that should occur after any attempted query */
1531
1532 sendquery_cleanup:
1533
1534         /* reset \g's output-to-filename trigger */
1535         if (pset.gfname)
1536         {
1537                 free(pset.gfname);
1538                 pset.gfname = NULL;
1539         }
1540
1541         /* reset \gx's expanded-mode flag */
1542         pset.g_expanded = false;
1543
1544         /* reset \gset trigger */
1545         if (pset.gset_prefix)
1546         {
1547                 free(pset.gset_prefix);
1548                 pset.gset_prefix = NULL;
1549         }
1550
1551         /* reset \gdesc trigger */
1552         pset.gdesc_flag = false;
1553
1554         /* reset \gexec trigger */
1555         pset.gexec_flag = false;
1556
1557         /* reset \crosstabview trigger */
1558         pset.crosstab_flag = false;
1559         for (i = 0; i < lengthof(pset.ctv_args); i++)
1560         {
1561                 pg_free(pset.ctv_args[i]);
1562                 pset.ctv_args[i] = NULL;
1563         }
1564
1565         return OK;
1566 }
1567
1568
1569 /*
1570  * DescribeQuery: describe the result columns of a query, without executing it
1571  *
1572  * Returns true if the operation executed successfully, false otherwise.
1573  *
1574  * If pset.timing is on, total query time (exclusive of result-printing) is
1575  * stored into *elapsed_msec.
1576  */
1577 static bool
1578 DescribeQuery(const char *query, double *elapsed_msec)
1579 {
1580         PGresult   *results;
1581         bool            OK;
1582         instr_time      before,
1583                                 after;
1584
1585         *elapsed_msec = 0;
1586
1587         if (pset.timing)
1588                 INSTR_TIME_SET_CURRENT(before);
1589
1590         /*
1591          * To parse the query but not execute it, we prepare it, using the unnamed
1592          * prepared statement.  This is invisible to psql users, since there's no
1593          * way to access the unnamed prepared statement from psql user space. The
1594          * next Parse or Query protocol message would overwrite the statement
1595          * anyway.  (So there's no great need to clear it when done, which is a
1596          * good thing because libpq provides no easy way to do that.)
1597          */
1598         results = PQprepare(pset.db, "", query, 0, NULL);
1599         if (PQresultStatus(results) != PGRES_COMMAND_OK)
1600         {
1601                 pg_log_info("%s", PQerrorMessage(pset.db));
1602                 SetResultVariables(results, false);
1603                 ClearOrSaveResult(results);
1604                 return false;
1605         }
1606         PQclear(results);
1607
1608         results = PQdescribePrepared(pset.db, "");
1609         OK = AcceptResult(results) &&
1610                 (PQresultStatus(results) == PGRES_COMMAND_OK);
1611         if (OK && results)
1612         {
1613                 if (PQnfields(results) > 0)
1614                 {
1615                         PQExpBufferData buf;
1616                         int                     i;
1617
1618                         initPQExpBuffer(&buf);
1619
1620                         printfPQExpBuffer(&buf,
1621                                                           "SELECT name AS \"%s\", pg_catalog.format_type(tp, tpm) AS \"%s\"\n"
1622                                                           "FROM (VALUES ",
1623                                                           gettext_noop("Column"),
1624                                                           gettext_noop("Type"));
1625
1626                         for (i = 0; i < PQnfields(results); i++)
1627                         {
1628                                 const char *name;
1629                                 char       *escname;
1630
1631                                 if (i > 0)
1632                                         appendPQExpBufferStr(&buf, ",");
1633
1634                                 name = PQfname(results, i);
1635                                 escname = PQescapeLiteral(pset.db, name, strlen(name));
1636
1637                                 if (escname == NULL)
1638                                 {
1639                                         pg_log_info("%s", PQerrorMessage(pset.db));
1640                                         PQclear(results);
1641                                         termPQExpBuffer(&buf);
1642                                         return false;
1643                                 }
1644
1645                                 appendPQExpBuffer(&buf, "(%s, '%u'::pg_catalog.oid, %d)",
1646                                                                   escname,
1647                                                                   PQftype(results, i),
1648                                                                   PQfmod(results, i));
1649
1650                                 PQfreemem(escname);
1651                         }
1652
1653                         appendPQExpBufferStr(&buf, ") s(name, tp, tpm)");
1654                         PQclear(results);
1655
1656                         results = PQexec(pset.db, buf.data);
1657                         OK = AcceptResult(results);
1658
1659                         if (pset.timing)
1660                         {
1661                                 INSTR_TIME_SET_CURRENT(after);
1662                                 INSTR_TIME_SUBTRACT(after, before);
1663                                 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1664                         }
1665
1666                         if (OK && results)
1667                                 OK = PrintQueryResults(results);
1668
1669                         termPQExpBuffer(&buf);
1670                 }
1671                 else
1672                         fprintf(pset.queryFout,
1673                                         _("The command has no result, or the result has no columns.\n"));
1674         }
1675
1676         SetResultVariables(results, OK);
1677         ClearOrSaveResult(results);
1678
1679         return OK;
1680 }
1681
1682
1683 /*
1684  * ExecQueryUsingCursor: run a SELECT-like query using a cursor
1685  *
1686  * This feature allows result sets larger than RAM to be dealt with.
1687  *
1688  * Returns true if the query executed successfully, false otherwise.
1689  *
1690  * If pset.timing is on, total query time (exclusive of result-printing) is
1691  * stored into *elapsed_msec.
1692  */
1693 static bool
1694 ExecQueryUsingCursor(const char *query, double *elapsed_msec)
1695 {
1696         bool            OK = true;
1697         PGresult   *results;
1698         PQExpBufferData buf;
1699         printQueryOpt my_popt = pset.popt;
1700         FILE       *fout;
1701         bool            is_pipe;
1702         bool            is_pager = false;
1703         bool            started_txn = false;
1704         int64           total_tuples = 0;
1705         int                     ntuples;
1706         int                     fetch_count;
1707         char            fetch_cmd[64];
1708         instr_time      before,
1709                                 after;
1710         int                     flush_error;
1711
1712         *elapsed_msec = 0;
1713
1714         /* initialize print options for partial table output */
1715         my_popt.topt.start_table = true;
1716         my_popt.topt.stop_table = false;
1717         my_popt.topt.prior_records = 0;
1718
1719         if (pset.timing)
1720                 INSTR_TIME_SET_CURRENT(before);
1721
1722         /* if we're not in a transaction, start one */
1723         if (PQtransactionStatus(pset.db) == PQTRANS_IDLE)
1724         {
1725                 results = PQexec(pset.db, "BEGIN");
1726                 OK = AcceptResult(results) &&
1727                         (PQresultStatus(results) == PGRES_COMMAND_OK);
1728                 ClearOrSaveResult(results);
1729                 if (!OK)
1730                         return false;
1731                 started_txn = true;
1732         }
1733
1734         /* Send DECLARE CURSOR */
1735         initPQExpBuffer(&buf);
1736         appendPQExpBuffer(&buf, "DECLARE _psql_cursor NO SCROLL CURSOR FOR\n%s",
1737                                           query);
1738
1739         results = PQexec(pset.db, buf.data);
1740         OK = AcceptResult(results) &&
1741                 (PQresultStatus(results) == PGRES_COMMAND_OK);
1742         if (!OK)
1743                 SetResultVariables(results, OK);
1744         ClearOrSaveResult(results);
1745         termPQExpBuffer(&buf);
1746         if (!OK)
1747                 goto cleanup;
1748
1749         if (pset.timing)
1750         {
1751                 INSTR_TIME_SET_CURRENT(after);
1752                 INSTR_TIME_SUBTRACT(after, before);
1753                 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1754         }
1755
1756         /*
1757          * In \gset mode, we force the fetch count to be 2, so that we will throw
1758          * the appropriate error if the query returns more than one row.
1759          */
1760         if (pset.gset_prefix)
1761                 fetch_count = 2;
1762         else
1763                 fetch_count = pset.fetch_count;
1764
1765         snprintf(fetch_cmd, sizeof(fetch_cmd),
1766                          "FETCH FORWARD %d FROM _psql_cursor",
1767                          fetch_count);
1768
1769         /* one-shot expanded output requested via \gx */
1770         if (pset.g_expanded)
1771                 my_popt.topt.expanded = 1;
1772
1773         /* prepare to write output to \g argument, if any */
1774         if (pset.gfname)
1775         {
1776                 if (!openQueryOutputFile(pset.gfname, &fout, &is_pipe))
1777                 {
1778                         OK = false;
1779                         goto cleanup;
1780                 }
1781                 if (is_pipe)
1782                         disable_sigpipe_trap();
1783         }
1784         else
1785         {
1786                 fout = pset.queryFout;
1787                 is_pipe = false;                /* doesn't matter */
1788         }
1789
1790         /* clear any pre-existing error indication on the output stream */
1791         clearerr(fout);
1792
1793         for (;;)
1794         {
1795                 if (pset.timing)
1796                         INSTR_TIME_SET_CURRENT(before);
1797
1798                 /* get fetch_count tuples at a time */
1799                 results = PQexec(pset.db, fetch_cmd);
1800
1801                 if (pset.timing)
1802                 {
1803                         INSTR_TIME_SET_CURRENT(after);
1804                         INSTR_TIME_SUBTRACT(after, before);
1805                         *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1806                 }
1807
1808                 if (PQresultStatus(results) != PGRES_TUPLES_OK)
1809                 {
1810                         /* shut down pager before printing error message */
1811                         if (is_pager)
1812                         {
1813                                 ClosePager(fout);
1814                                 is_pager = false;
1815                         }
1816
1817                         OK = AcceptResult(results);
1818                         Assert(!OK);
1819                         SetResultVariables(results, OK);
1820                         ClearOrSaveResult(results);
1821                         break;
1822                 }
1823
1824                 if (pset.gset_prefix)
1825                 {
1826                         /* StoreQueryTuple will complain if not exactly one row */
1827                         OK = StoreQueryTuple(results);
1828                         ClearOrSaveResult(results);
1829                         break;
1830                 }
1831
1832                 /*
1833                  * Note we do not deal with \gdesc, \gexec or \crosstabview modes here
1834                  */
1835
1836                 ntuples = PQntuples(results);
1837                 total_tuples += ntuples;
1838
1839                 if (ntuples < fetch_count)
1840                 {
1841                         /* this is the last result set, so allow footer decoration */
1842                         my_popt.topt.stop_table = true;
1843                 }
1844                 else if (fout == stdout && !is_pager)
1845                 {
1846                         /*
1847                          * If query requires multiple result sets, hack to ensure that
1848                          * only one pager instance is used for the whole mess
1849                          */
1850                         fout = PageOutput(INT_MAX, &(my_popt.topt));
1851                         is_pager = true;
1852                 }
1853
1854                 printQuery(results, &my_popt, fout, is_pager, pset.logfile);
1855
1856                 ClearOrSaveResult(results);
1857
1858                 /* after the first result set, disallow header decoration */
1859                 my_popt.topt.start_table = false;
1860                 my_popt.topt.prior_records += ntuples;
1861
1862                 /*
1863                  * Make sure to flush the output stream, so intermediate results are
1864                  * visible to the client immediately.  We check the results because if
1865                  * the pager dies/exits/etc, there's no sense throwing more data at
1866                  * it.
1867                  */
1868                 flush_error = fflush(fout);
1869
1870                 /*
1871                  * Check if we are at the end, if a cancel was pressed, or if there
1872                  * were any errors either trying to flush out the results, or more
1873                  * generally on the output stream at all.  If we hit any errors
1874                  * writing things to the stream, we presume $PAGER has disappeared and
1875                  * stop bothering to pull down more data.
1876                  */
1877                 if (ntuples < fetch_count || cancel_pressed || flush_error ||
1878                         ferror(fout))
1879                         break;
1880         }
1881
1882         if (pset.gfname)
1883         {
1884                 /* close \g argument file/pipe */
1885                 if (is_pipe)
1886                 {
1887                         pclose(fout);
1888                         restore_sigpipe_trap();
1889                 }
1890                 else
1891                         fclose(fout);
1892         }
1893         else if (is_pager)
1894         {
1895                 /* close transient pager */
1896                 ClosePager(fout);
1897         }
1898
1899         if (OK)
1900         {
1901                 /*
1902                  * We don't have a PGresult here, and even if we did it wouldn't have
1903                  * the right row count, so fake SetResultVariables().  In error cases,
1904                  * we already set the result variables above.
1905                  */
1906                 char            buf[32];
1907
1908                 SetVariable(pset.vars, "ERROR", "false");
1909                 SetVariable(pset.vars, "SQLSTATE", "00000");
1910                 snprintf(buf, sizeof(buf), INT64_FORMAT, total_tuples);
1911                 SetVariable(pset.vars, "ROW_COUNT", buf);
1912         }
1913
1914 cleanup:
1915         if (pset.timing)
1916                 INSTR_TIME_SET_CURRENT(before);
1917
1918         /*
1919          * We try to close the cursor on either success or failure, but on failure
1920          * ignore the result (it's probably just a bleat about being in an aborted
1921          * transaction)
1922          */
1923         results = PQexec(pset.db, "CLOSE _psql_cursor");
1924         if (OK)
1925         {
1926                 OK = AcceptResult(results) &&
1927                         (PQresultStatus(results) == PGRES_COMMAND_OK);
1928                 ClearOrSaveResult(results);
1929         }
1930         else
1931                 PQclear(results);
1932
1933         if (started_txn)
1934         {
1935                 results = PQexec(pset.db, OK ? "COMMIT" : "ROLLBACK");
1936                 OK &= AcceptResult(results) &&
1937                         (PQresultStatus(results) == PGRES_COMMAND_OK);
1938                 ClearOrSaveResult(results);
1939         }
1940
1941         if (pset.timing)
1942         {
1943                 INSTR_TIME_SET_CURRENT(after);
1944                 INSTR_TIME_SUBTRACT(after, before);
1945                 *elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
1946         }
1947
1948         return OK;
1949 }
1950
1951
1952 /*
1953  * Advance the given char pointer over white space and SQL comments.
1954  */
1955 static const char *
1956 skip_white_space(const char *query)
1957 {
1958         int                     cnestlevel = 0; /* slash-star comment nest level */
1959
1960         while (*query)
1961         {
1962                 int                     mblen = PQmblen(query, pset.encoding);
1963
1964                 /*
1965                  * Note: we assume the encoding is a superset of ASCII, so that for
1966                  * example "query[0] == '/'" is meaningful.  However, we do NOT assume
1967                  * that the second and subsequent bytes of a multibyte character
1968                  * couldn't look like ASCII characters; so it is critical to advance
1969                  * by mblen, not 1, whenever we haven't exactly identified the
1970                  * character we are skipping over.
1971                  */
1972                 if (isspace((unsigned char) *query))
1973                         query += mblen;
1974                 else if (query[0] == '/' && query[1] == '*')
1975                 {
1976                         cnestlevel++;
1977                         query += 2;
1978                 }
1979                 else if (cnestlevel > 0 && query[0] == '*' && query[1] == '/')
1980                 {
1981                         cnestlevel--;
1982                         query += 2;
1983                 }
1984                 else if (cnestlevel == 0 && query[0] == '-' && query[1] == '-')
1985                 {
1986                         query += 2;
1987
1988                         /*
1989                          * We have to skip to end of line since any slash-star inside the
1990                          * -- comment does NOT start a slash-star comment.
1991                          */
1992                         while (*query)
1993                         {
1994                                 if (*query == '\n')
1995                                 {
1996                                         query++;
1997                                         break;
1998                                 }
1999                                 query += PQmblen(query, pset.encoding);
2000                         }
2001                 }
2002                 else if (cnestlevel > 0)
2003                         query += mblen;
2004                 else
2005                         break;                          /* found first token */
2006         }
2007
2008         return query;
2009 }
2010
2011
2012 /*
2013  * Check whether a command is one of those for which we should NOT start
2014  * a new transaction block (ie, send a preceding BEGIN).
2015  *
2016  * These include the transaction control statements themselves, plus
2017  * certain statements that the backend disallows inside transaction blocks.
2018  */
2019 static bool
2020 command_no_begin(const char *query)
2021 {
2022         int                     wordlen;
2023
2024         /*
2025          * First we must advance over any whitespace and comments.
2026          */
2027         query = skip_white_space(query);
2028
2029         /*
2030          * Check word length (since "beginx" is not "begin").
2031          */
2032         wordlen = 0;
2033         while (isalpha((unsigned char) query[wordlen]))
2034                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2035
2036         /*
2037          * Transaction control commands.  These should include every keyword that
2038          * gives rise to a TransactionStmt in the backend grammar, except for the
2039          * savepoint-related commands.
2040          *
2041          * (We assume that START must be START TRANSACTION, since there is
2042          * presently no other "START foo" command.)
2043          */
2044         if (wordlen == 5 && pg_strncasecmp(query, "abort", 5) == 0)
2045                 return true;
2046         if (wordlen == 5 && pg_strncasecmp(query, "begin", 5) == 0)
2047                 return true;
2048         if (wordlen == 5 && pg_strncasecmp(query, "start", 5) == 0)
2049                 return true;
2050         if (wordlen == 6 && pg_strncasecmp(query, "commit", 6) == 0)
2051                 return true;
2052         if (wordlen == 3 && pg_strncasecmp(query, "end", 3) == 0)
2053                 return true;
2054         if (wordlen == 8 && pg_strncasecmp(query, "rollback", 8) == 0)
2055                 return true;
2056         if (wordlen == 7 && pg_strncasecmp(query, "prepare", 7) == 0)
2057         {
2058                 /* PREPARE TRANSACTION is a TC command, PREPARE foo is not */
2059                 query += wordlen;
2060
2061                 query = skip_white_space(query);
2062
2063                 wordlen = 0;
2064                 while (isalpha((unsigned char) query[wordlen]))
2065                         wordlen += PQmblen(&query[wordlen], pset.encoding);
2066
2067                 if (wordlen == 11 && pg_strncasecmp(query, "transaction", 11) == 0)
2068                         return true;
2069                 return false;
2070         }
2071
2072         /*
2073          * Commands not allowed within transactions.  The statements checked for
2074          * here should be exactly those that call PreventInTransactionBlock() in
2075          * the backend.
2076          */
2077         if (wordlen == 6 && pg_strncasecmp(query, "vacuum", 6) == 0)
2078                 return true;
2079         if (wordlen == 7 && pg_strncasecmp(query, "cluster", 7) == 0)
2080         {
2081                 /* CLUSTER with any arguments is allowed in transactions */
2082                 query += wordlen;
2083
2084                 query = skip_white_space(query);
2085
2086                 if (isalpha((unsigned char) query[0]))
2087                         return false;           /* has additional words */
2088                 return true;                    /* it's CLUSTER without arguments */
2089         }
2090
2091         if (wordlen == 6 && pg_strncasecmp(query, "create", 6) == 0)
2092         {
2093                 query += wordlen;
2094
2095                 query = skip_white_space(query);
2096
2097                 wordlen = 0;
2098                 while (isalpha((unsigned char) query[wordlen]))
2099                         wordlen += PQmblen(&query[wordlen], pset.encoding);
2100
2101                 if (wordlen == 8 && pg_strncasecmp(query, "database", 8) == 0)
2102                         return true;
2103                 if (wordlen == 10 && pg_strncasecmp(query, "tablespace", 10) == 0)
2104                         return true;
2105
2106                 /* CREATE [UNIQUE] INDEX CONCURRENTLY isn't allowed in xacts */
2107                 if (wordlen == 6 && pg_strncasecmp(query, "unique", 6) == 0)
2108                 {
2109                         query += wordlen;
2110
2111                         query = skip_white_space(query);
2112
2113                         wordlen = 0;
2114                         while (isalpha((unsigned char) query[wordlen]))
2115                                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2116                 }
2117
2118                 if (wordlen == 5 && pg_strncasecmp(query, "index", 5) == 0)
2119                 {
2120                         query += wordlen;
2121
2122                         query = skip_white_space(query);
2123
2124                         wordlen = 0;
2125                         while (isalpha((unsigned char) query[wordlen]))
2126                                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2127
2128                         if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2129                                 return true;
2130                 }
2131
2132                 return false;
2133         }
2134
2135         if (wordlen == 5 && pg_strncasecmp(query, "alter", 5) == 0)
2136         {
2137                 query += wordlen;
2138
2139                 query = skip_white_space(query);
2140
2141                 wordlen = 0;
2142                 while (isalpha((unsigned char) query[wordlen]))
2143                         wordlen += PQmblen(&query[wordlen], pset.encoding);
2144
2145                 /* ALTER SYSTEM isn't allowed in xacts */
2146                 if (wordlen == 6 && pg_strncasecmp(query, "system", 6) == 0)
2147                         return true;
2148
2149                 return false;
2150         }
2151
2152         /*
2153          * Note: these tests will match DROP SYSTEM and REINDEX TABLESPACE, which
2154          * aren't really valid commands so we don't care much. The other four
2155          * possible matches are correct.
2156          */
2157         if ((wordlen == 4 && pg_strncasecmp(query, "drop", 4) == 0) ||
2158                 (wordlen == 7 && pg_strncasecmp(query, "reindex", 7) == 0))
2159         {
2160                 query += wordlen;
2161
2162                 query = skip_white_space(query);
2163
2164                 wordlen = 0;
2165                 while (isalpha((unsigned char) query[wordlen]))
2166                         wordlen += PQmblen(&query[wordlen], pset.encoding);
2167
2168                 if (wordlen == 8 && pg_strncasecmp(query, "database", 8) == 0)
2169                         return true;
2170                 if (wordlen == 6 && pg_strncasecmp(query, "system", 6) == 0)
2171                         return true;
2172                 if (wordlen == 10 && pg_strncasecmp(query, "tablespace", 10) == 0)
2173                         return true;
2174                 if (wordlen == 5 && (pg_strncasecmp(query, "index", 5) == 0 ||
2175                                                          pg_strncasecmp(query, "table", 5) == 0))
2176                 {
2177                         query += wordlen;
2178                         query = skip_white_space(query);
2179                         wordlen = 0;
2180                         while (isalpha((unsigned char) query[wordlen]))
2181                                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2182
2183                         /*
2184                          * REINDEX [ TABLE | INDEX ] CONCURRENTLY are not allowed in
2185                          * xacts.
2186                          */
2187                         if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2188                                 return true;
2189                 }
2190
2191                 /* DROP INDEX CONCURRENTLY isn't allowed in xacts */
2192                 if (wordlen == 5 && pg_strncasecmp(query, "index", 5) == 0)
2193                 {
2194                         query += wordlen;
2195
2196                         query = skip_white_space(query);
2197
2198                         wordlen = 0;
2199                         while (isalpha((unsigned char) query[wordlen]))
2200                                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2201
2202                         if (wordlen == 12 && pg_strncasecmp(query, "concurrently", 12) == 0)
2203                                 return true;
2204
2205                         return false;
2206                 }
2207
2208                 return false;
2209         }
2210
2211         /* DISCARD ALL isn't allowed in xacts, but other variants are allowed. */
2212         if (wordlen == 7 && pg_strncasecmp(query, "discard", 7) == 0)
2213         {
2214                 query += wordlen;
2215
2216                 query = skip_white_space(query);
2217
2218                 wordlen = 0;
2219                 while (isalpha((unsigned char) query[wordlen]))
2220                         wordlen += PQmblen(&query[wordlen], pset.encoding);
2221
2222                 if (wordlen == 3 && pg_strncasecmp(query, "all", 3) == 0)
2223                         return true;
2224                 return false;
2225         }
2226
2227         return false;
2228 }
2229
2230
2231 /*
2232  * Check whether the specified command is a SELECT (or VALUES).
2233  */
2234 static bool
2235 is_select_command(const char *query)
2236 {
2237         int                     wordlen;
2238
2239         /*
2240          * First advance over any whitespace, comments and left parentheses.
2241          */
2242         for (;;)
2243         {
2244                 query = skip_white_space(query);
2245                 if (query[0] == '(')
2246                         query++;
2247                 else
2248                         break;
2249         }
2250
2251         /*
2252          * Check word length (since "selectx" is not "select").
2253          */
2254         wordlen = 0;
2255         while (isalpha((unsigned char) query[wordlen]))
2256                 wordlen += PQmblen(&query[wordlen], pset.encoding);
2257
2258         if (wordlen == 6 && pg_strncasecmp(query, "select", 6) == 0)
2259                 return true;
2260
2261         if (wordlen == 6 && pg_strncasecmp(query, "values", 6) == 0)
2262                 return true;
2263
2264         return false;
2265 }
2266
2267
2268 /*
2269  * Test if the current user is a database superuser.
2270  *
2271  * Note: this will correctly detect superuserness only with a protocol-3.0
2272  * or newer backend; otherwise it will always say "false".
2273  */
2274 bool
2275 is_superuser(void)
2276 {
2277         const char *val;
2278
2279         if (!pset.db)
2280                 return false;
2281
2282         val = PQparameterStatus(pset.db, "is_superuser");
2283
2284         if (val && strcmp(val, "on") == 0)
2285                 return true;
2286
2287         return false;
2288 }
2289
2290
2291 /*
2292  * Test if the current session uses standard string literals.
2293  *
2294  * Note: With a pre-protocol-3.0 connection this will always say "false",
2295  * which should be the right answer.
2296  */
2297 bool
2298 standard_strings(void)
2299 {
2300         const char *val;
2301
2302         if (!pset.db)
2303                 return false;
2304
2305         val = PQparameterStatus(pset.db, "standard_conforming_strings");
2306
2307         if (val && strcmp(val, "on") == 0)
2308                 return true;
2309
2310         return false;
2311 }
2312
2313
2314 /*
2315  * Return the session user of the current connection.
2316  *
2317  * Note: this will correctly detect the session user only with a
2318  * protocol-3.0 or newer backend; otherwise it will return the
2319  * connection user.
2320  */
2321 const char *
2322 session_username(void)
2323 {
2324         const char *val;
2325
2326         if (!pset.db)
2327                 return NULL;
2328
2329         val = PQparameterStatus(pset.db, "session_authorization");
2330         if (val)
2331                 return val;
2332         else
2333                 return PQuser(pset.db);
2334 }
2335
2336
2337 /* expand_tilde
2338  *
2339  * substitute '~' with HOME or '~username' with username's home dir
2340  *
2341  */
2342 void
2343 expand_tilde(char **filename)
2344 {
2345         if (!filename || !(*filename))
2346                 return;
2347
2348         /*
2349          * WIN32 doesn't use tilde expansion for file names. Also, it uses tilde
2350          * for short versions of long file names, though the tilde is usually
2351          * toward the end, not at the beginning.
2352          */
2353 #ifndef WIN32
2354
2355         /* try tilde expansion */
2356         if (**filename == '~')
2357         {
2358                 char       *fn;
2359                 char            oldp,
2360                                    *p;
2361                 struct passwd *pw;
2362                 char            home[MAXPGPATH];
2363
2364                 fn = *filename;
2365                 *home = '\0';
2366
2367                 p = fn + 1;
2368                 while (*p != '/' && *p != '\0')
2369                         p++;
2370
2371                 oldp = *p;
2372                 *p = '\0';
2373
2374                 if (*(fn + 1) == '\0')
2375                         get_home_path(home);    /* ~ or ~/ only */
2376                 else if ((pw = getpwnam(fn + 1)) != NULL)
2377                         strlcpy(home, pw->pw_dir, sizeof(home));        /* ~user */
2378
2379                 *p = oldp;
2380                 if (strlen(home) != 0)
2381                 {
2382                         char       *newfn;
2383
2384                         newfn = psprintf("%s%s", home, p);
2385                         free(fn);
2386                         *filename = newfn;
2387                 }
2388         }
2389 #endif
2390
2391         return;
2392 }
2393
2394 /*
2395  * Checks if connection string starts with either of the valid URI prefix
2396  * designators.
2397  *
2398  * Returns the URI prefix length, 0 if the string doesn't contain a URI prefix.
2399  *
2400  * XXX This is a duplicate of the eponymous libpq function.
2401  */
2402 static int
2403 uri_prefix_length(const char *connstr)
2404 {
2405         /* The connection URI must start with either of the following designators: */
2406         static const char uri_designator[] = "postgresql://";
2407         static const char short_uri_designator[] = "postgres://";
2408
2409         if (strncmp(connstr, uri_designator,
2410                                 sizeof(uri_designator) - 1) == 0)
2411                 return sizeof(uri_designator) - 1;
2412
2413         if (strncmp(connstr, short_uri_designator,
2414                                 sizeof(short_uri_designator) - 1) == 0)
2415                 return sizeof(short_uri_designator) - 1;
2416
2417         return 0;
2418 }
2419
2420 /*
2421  * Recognized connection string either starts with a valid URI prefix or
2422  * contains a "=" in it.
2423  *
2424  * Must be consistent with parse_connection_string: anything for which this
2425  * returns true should at least look like it's parseable by that routine.
2426  *
2427  * XXX This is a duplicate of the eponymous libpq function.
2428  */
2429 bool
2430 recognized_connection_string(const char *connstr)
2431 {
2432         return uri_prefix_length(connstr) != 0 || strchr(connstr, '=') != NULL;
2433 }