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