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

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