LCOV - code coverage report
Current view: top level - src/bin/scripts - common.c (source / functions) Hit Total Coverage
Test: PostgreSQL 13devel Lines: 138 184 75.0 %
Date: 2019-09-22 08:06:49 Functions: 14 16 87.5 %
Legend: Lines: hit not hit

          Line data    Source code
       1             : /*-------------------------------------------------------------------------
       2             :  *
       3             :  *  common.c
       4             :  *      Common support routines for bin/scripts/
       5             :  *
       6             :  *
       7             :  * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
       8             :  * Portions Copyright (c) 1994, Regents of the University of California
       9             :  *
      10             :  * src/bin/scripts/common.c
      11             :  *
      12             :  *-------------------------------------------------------------------------
      13             :  */
      14             : 
      15             : #include "postgres_fe.h"
      16             : 
      17             : #include <signal.h>
      18             : #include <unistd.h>
      19             : 
      20             : #include "common.h"
      21             : #include "common/logging.h"
      22             : #include "fe_utils/connect.h"
      23             : #include "fe_utils/string_utils.h"
      24             : 
      25             : #define ERRCODE_UNDEFINED_TABLE  "42P01"
      26             : 
      27             : 
      28             : static PGcancel *volatile cancelConn = NULL;
      29             : bool        CancelRequested = false;
      30             : 
      31             : #ifdef WIN32
      32             : static CRITICAL_SECTION cancelConnLock;
      33             : #endif
      34             : 
      35             : /*
      36             :  * Provide strictly harmonized handling of --help and --version
      37             :  * options.
      38             :  */
      39             : void
      40         238 : handle_help_version_opts(int argc, char *argv[],
      41             :                          const char *fixed_progname, help_handler hlp)
      42             : {
      43         238 :     if (argc > 1)
      44             :     {
      45         234 :         if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
      46             :         {
      47          16 :             hlp(get_progname(argv[0]));
      48          16 :             exit(0);
      49             :         }
      50         218 :         if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
      51             :         {
      52          16 :             printf("%s (PostgreSQL) " PG_VERSION "\n", fixed_progname);
      53          16 :             exit(0);
      54             :         }
      55             :     }
      56         206 : }
      57             : 
      58             : 
      59             : /*
      60             :  * Make a database connection with the given parameters.
      61             :  *
      62             :  * An interactive password prompt is automatically issued if needed and
      63             :  * allowed by prompt_password.
      64             :  *
      65             :  * If allow_password_reuse is true, we will try to re-use any password
      66             :  * given during previous calls to this routine.  (Callers should not pass
      67             :  * allow_password_reuse=true unless reconnecting to the same database+user
      68             :  * as before, else we might create password exposure hazards.)
      69             :  */
      70             : PGconn *
      71         298 : connectDatabase(const char *dbname, const char *pghost,
      72             :                 const char *pgport, const char *pguser,
      73             :                 enum trivalue prompt_password, const char *progname,
      74             :                 bool echo, bool fail_ok, bool allow_password_reuse)
      75             : {
      76             :     PGconn     *conn;
      77             :     bool        new_pass;
      78             :     static bool have_password = false;
      79             :     static char password[100];
      80             : 
      81         298 :     if (!allow_password_reuse)
      82         172 :         have_password = false;
      83             : 
      84         298 :     if (!have_password && prompt_password == TRI_YES)
      85             :     {
      86           0 :         simple_prompt("Password: ", password, sizeof(password), false);
      87           0 :         have_password = true;
      88             :     }
      89             : 
      90             :     /*
      91             :      * Start the connection.  Loop until we have a password if requested by
      92             :      * backend.
      93             :      */
      94             :     do
      95             :     {
      96             :         const char *keywords[7];
      97             :         const char *values[7];
      98             : 
      99         298 :         keywords[0] = "host";
     100         298 :         values[0] = pghost;
     101         298 :         keywords[1] = "port";
     102         298 :         values[1] = pgport;
     103         298 :         keywords[2] = "user";
     104         298 :         values[2] = pguser;
     105         298 :         keywords[3] = "password";
     106         298 :         values[3] = have_password ? password : NULL;
     107         298 :         keywords[4] = "dbname";
     108         298 :         values[4] = dbname;
     109         298 :         keywords[5] = "fallback_application_name";
     110         298 :         values[5] = progname;
     111         298 :         keywords[6] = NULL;
     112         298 :         values[6] = NULL;
     113             : 
     114         298 :         new_pass = false;
     115         298 :         conn = PQconnectdbParams(keywords, values, true);
     116             : 
     117         298 :         if (!conn)
     118             :         {
     119           0 :             pg_log_error("could not connect to database %s: out of memory",
     120             :                          dbname);
     121           0 :             exit(1);
     122             :         }
     123             : 
     124             :         /*
     125             :          * No luck?  Trying asking (again) for a password.
     126             :          */
     127         298 :         if (PQstatus(conn) == CONNECTION_BAD &&
     128           0 :             PQconnectionNeedsPassword(conn) &&
     129             :             prompt_password != TRI_NO)
     130             :         {
     131           0 :             PQfinish(conn);
     132           0 :             simple_prompt("Password: ", password, sizeof(password), false);
     133           0 :             have_password = true;
     134           0 :             new_pass = true;
     135             :         }
     136         298 :     } while (new_pass);
     137             : 
     138             :     /* check to see that the backend connection was successfully made */
     139         298 :     if (PQstatus(conn) == CONNECTION_BAD)
     140             :     {
     141           0 :         if (fail_ok)
     142             :         {
     143           0 :             PQfinish(conn);
     144           0 :             return NULL;
     145             :         }
     146           0 :         pg_log_error("could not connect to database %s: %s",
     147             :                      dbname, PQerrorMessage(conn));
     148           0 :         exit(1);
     149             :     }
     150             : 
     151         298 :     PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL, echo));
     152             : 
     153         298 :     return conn;
     154             : }
     155             : 
     156             : /*
     157             :  * Try to connect to the appropriate maintenance database.
     158             :  */
     159             : PGconn *
     160          64 : connectMaintenanceDatabase(const char *maintenance_db,
     161             :                            const char *pghost, const char *pgport,
     162             :                            const char *pguser, enum trivalue prompt_password,
     163             :                            const char *progname, bool echo)
     164             : {
     165             :     PGconn     *conn;
     166             : 
     167             :     /* If a maintenance database name was specified, just connect to it. */
     168          64 :     if (maintenance_db)
     169           0 :         return connectDatabase(maintenance_db, pghost, pgport, pguser,
     170             :                                prompt_password, progname, echo, false, false);
     171             : 
     172             :     /* Otherwise, try postgres first and then template1. */
     173          64 :     conn = connectDatabase("postgres", pghost, pgport, pguser, prompt_password,
     174             :                            progname, echo, true, false);
     175          64 :     if (!conn)
     176           0 :         conn = connectDatabase("template1", pghost, pgport, pguser,
     177             :                                prompt_password, progname, echo, false, false);
     178             : 
     179          64 :     return conn;
     180             : }
     181             : 
     182             : /*
     183             :  * Disconnect the given connection, canceling any statement if one is active.
     184             :  */
     185             : void
     186         174 : disconnectDatabase(PGconn *conn)
     187             : {
     188             :     char        errbuf[256];
     189             : 
     190             :     Assert(conn != NULL);
     191             : 
     192         174 :     if (PQtransactionStatus(conn) == PQTRANS_ACTIVE)
     193             :     {
     194             :         PGcancel   *cancel;
     195             : 
     196           0 :         if ((cancel = PQgetCancel(conn)))
     197             :         {
     198           0 :             (void) PQcancel(cancel, errbuf, sizeof(errbuf));
     199           0 :             PQfreeCancel(cancel);
     200             :         }
     201             :     }
     202             : 
     203         174 :     PQfinish(conn);
     204         174 : }
     205             : 
     206             : /*
     207             :  * Run a query, return the results, exit program on failure.
     208             :  */
     209             : PGresult *
     210         682 : executeQuery(PGconn *conn, const char *query, bool echo)
     211             : {
     212             :     PGresult   *res;
     213             : 
     214         682 :     if (echo)
     215          94 :         printf("%s\n", query);
     216             : 
     217         682 :     res = PQexec(conn, query);
     218        1364 :     if (!res ||
     219         682 :         PQresultStatus(res) != PGRES_TUPLES_OK)
     220             :     {
     221           4 :         pg_log_error("query failed: %s", PQerrorMessage(conn));
     222           4 :         pg_log_info("query was: %s", query);
     223           4 :         PQfinish(conn);
     224           4 :         exit(1);
     225             :     }
     226             : 
     227         678 :     return res;
     228             : }
     229             : 
     230             : 
     231             : /*
     232             :  * As above for a SQL command (which returns nothing).
     233             :  */
     234             : void
     235         234 : executeCommand(PGconn *conn, const char *query, bool echo)
     236             : {
     237             :     PGresult   *res;
     238             : 
     239         234 :     if (echo)
     240          16 :         printf("%s\n", query);
     241             : 
     242         234 :     res = PQexec(conn, query);
     243         468 :     if (!res ||
     244         234 :         PQresultStatus(res) != PGRES_COMMAND_OK)
     245             :     {
     246           0 :         pg_log_error("query failed: %s", PQerrorMessage(conn));
     247           0 :         pg_log_info("query was: %s", query);
     248           0 :         PQfinish(conn);
     249           0 :         exit(1);
     250             :     }
     251             : 
     252         234 :     PQclear(res);
     253         234 : }
     254             : 
     255             : 
     256             : /*
     257             :  * As above for a SQL maintenance command (returns command success).
     258             :  * Command is executed with a cancel handler set, so Ctrl-C can
     259             :  * interrupt it.
     260             :  */
     261             : bool
     262          26 : executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
     263             : {
     264             :     PGresult   *res;
     265             :     bool        r;
     266             : 
     267          26 :     if (echo)
     268          16 :         printf("%s\n", query);
     269             : 
     270          26 :     SetCancelConn(conn);
     271          26 :     res = PQexec(conn, query);
     272          26 :     ResetCancelConn();
     273             : 
     274          26 :     r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);
     275             : 
     276          26 :     if (res)
     277          26 :         PQclear(res);
     278             : 
     279          26 :     return r;
     280             : }
     281             : 
     282             : /*
     283             :  * Consume all the results generated for the given connection until
     284             :  * nothing remains.  If at least one error is encountered, return false.
     285             :  * Note that this will block if the connection is busy.
     286             :  */
     287             : bool
     288         174 : consumeQueryResult(PGconn *conn)
     289             : {
     290         174 :     bool        ok = true;
     291             :     PGresult   *result;
     292             : 
     293         174 :     SetCancelConn(conn);
     294         524 :     while ((result = PQgetResult(conn)) != NULL)
     295             :     {
     296         176 :         if (!processQueryResult(conn, result))
     297           6 :             ok = false;
     298             :     }
     299         174 :     ResetCancelConn();
     300         174 :     return ok;
     301             : }
     302             : 
     303             : /*
     304             :  * Process (and delete) a query result.  Returns true if there's no error,
     305             :  * false otherwise -- but errors about trying to work on a missing relation
     306             :  * are reported and subsequently ignored.
     307             :  */
     308             : bool
     309        9626 : processQueryResult(PGconn *conn, PGresult *result)
     310             : {
     311             :     /*
     312             :      * If it's an error, report it.  Errors about a missing table are harmless
     313             :      * so we continue processing; but die for other errors.
     314             :      */
     315        9626 :     if (PQresultStatus(result) != PGRES_COMMAND_OK)
     316             :     {
     317           6 :         char       *sqlState = PQresultErrorField(result, PG_DIAG_SQLSTATE);
     318             : 
     319           6 :         pg_log_error("processing of database \"%s\" failed: %s",
     320             :                      PQdb(conn), PQerrorMessage(conn));
     321             : 
     322           6 :         if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) != 0)
     323             :         {
     324           6 :             PQclear(result);
     325           6 :             return false;
     326             :         }
     327             :     }
     328             : 
     329        9620 :     PQclear(result);
     330        9620 :     return true;
     331             : }
     332             : 
     333             : 
     334             : /*
     335             :  * Split TABLE[(COLUMNS)] into TABLE and [(COLUMNS)] portions.  When you
     336             :  * finish using them, pg_free(*table).  *columns is a pointer into "spec",
     337             :  * possibly to its NUL terminator.
     338             :  */
     339             : void
     340          82 : splitTableColumnsSpec(const char *spec, int encoding,
     341             :                       char **table, const char **columns)
     342             : {
     343          82 :     bool        inquotes = false;
     344          82 :     const char *cp = spec;
     345             : 
     346             :     /*
     347             :      * Find the first '(' not identifier-quoted.  Based on
     348             :      * dequote_downcase_identifier().
     349             :      */
     350        1460 :     while (*cp && (*cp != '(' || inquotes))
     351             :     {
     352        1296 :         if (*cp == '"')
     353             :         {
     354           6 :             if (inquotes && cp[1] == '"')
     355           2 :                 cp++;           /* pair does not affect quoting */
     356             :             else
     357           4 :                 inquotes = !inquotes;
     358           6 :             cp++;
     359             :         }
     360             :         else
     361        1290 :             cp += PQmblen(cp, encoding);
     362             :     }
     363          82 :     *table = pg_strdup(spec);
     364          82 :     (*table)[cp - spec] = '\0'; /* no strndup */
     365          82 :     *columns = cp;
     366          82 : }
     367             : 
     368             : /*
     369             :  * Break apart TABLE[(COLUMNS)] of "spec".  With the reset_val of search_path
     370             :  * in effect, have regclassin() interpret the TABLE portion.  Append to "buf"
     371             :  * the qualified name of TABLE, followed by any (COLUMNS).  Exit on failure.
     372             :  * We use this to interpret --table=foo under the search path psql would get,
     373             :  * in advance of "ANALYZE public.foo" under the always-secure search path.
     374             :  */
     375             : void
     376          62 : appendQualifiedRelation(PQExpBuffer buf, const char *spec,
     377             :                         PGconn *conn, bool echo)
     378             : {
     379             :     char       *table;
     380             :     const char *columns;
     381             :     PQExpBufferData sql;
     382             :     PGresult   *res;
     383             :     int         ntups;
     384             : 
     385          62 :     splitTableColumnsSpec(spec, PQclientEncoding(conn), &table, &columns);
     386             : 
     387             :     /*
     388             :      * Query must remain ABSOLUTELY devoid of unqualified names.  This would
     389             :      * be unnecessary given a regclassin() variant taking a search_path
     390             :      * argument.
     391             :      */
     392          62 :     initPQExpBuffer(&sql);
     393          62 :     appendPQExpBufferStr(&sql,
     394             :                          "SELECT c.relname, ns.nspname\n"
     395             :                          " FROM pg_catalog.pg_class c,"
     396             :                          " pg_catalog.pg_namespace ns\n"
     397             :                          " WHERE c.relnamespace OPERATOR(pg_catalog.=) ns.oid\n"
     398             :                          "  AND c.oid OPERATOR(pg_catalog.=) ");
     399          62 :     appendStringLiteralConn(&sql, table, conn);
     400          62 :     appendPQExpBufferStr(&sql, "::pg_catalog.regclass;");
     401             : 
     402          62 :     executeCommand(conn, "RESET search_path;", echo);
     403             : 
     404             :     /*
     405             :      * One row is a typical result, as is a nonexistent relation ERROR.
     406             :      * regclassin() unconditionally accepts all-digits input as an OID; if no
     407             :      * relation has that OID; this query returns no rows.  Catalog corruption
     408             :      * might elicit other row counts.
     409             :      */
     410          62 :     res = executeQuery(conn, sql.data, echo);
     411          60 :     ntups = PQntuples(res);
     412          60 :     if (ntups != 1)
     413             :     {
     414           0 :         pg_log_error(ngettext("query returned %d row instead of one: %s",
     415             :                               "query returned %d rows instead of one: %s",
     416             :                               ntups),
     417             :                      ntups, sql.data);
     418           0 :         PQfinish(conn);
     419           0 :         exit(1);
     420             :     }
     421          60 :     appendPQExpBufferStr(buf,
     422          60 :                          fmtQualifiedId(PQgetvalue(res, 0, 1),
     423          60 :                                         PQgetvalue(res, 0, 0)));
     424          60 :     appendPQExpBufferStr(buf, columns);
     425          60 :     PQclear(res);
     426          60 :     termPQExpBuffer(&sql);
     427          60 :     pg_free(table);
     428             : 
     429          60 :     PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL, echo));
     430          60 : }
     431             : 
     432             : 
     433             : /*
     434             :  * Check yes/no answer in a localized way.  1=yes, 0=no, -1=neither.
     435             :  */
     436             : 
     437             : /* translator: abbreviation for "yes" */
     438             : #define PG_YESLETTER gettext_noop("y")
     439             : /* translator: abbreviation for "no" */
     440             : #define PG_NOLETTER gettext_noop("n")
     441             : 
     442             : bool
     443           0 : yesno_prompt(const char *question)
     444             : {
     445             :     char        prompt[256];
     446             : 
     447             :     /*------
     448             :        translator: This is a question followed by the translated options for
     449             :        "yes" and "no". */
     450           0 :     snprintf(prompt, sizeof(prompt), _("%s (%s/%s) "),
     451             :              _(question), _(PG_YESLETTER), _(PG_NOLETTER));
     452             : 
     453             :     for (;;)
     454           0 :     {
     455             :         char        resp[10];
     456             : 
     457           0 :         simple_prompt(prompt, resp, sizeof(resp), true);
     458             : 
     459           0 :         if (strcmp(resp, _(PG_YESLETTER)) == 0)
     460           0 :             return true;
     461           0 :         if (strcmp(resp, _(PG_NOLETTER)) == 0)
     462           0 :             return false;
     463             : 
     464           0 :         printf(_("Please answer \"%s\" or \"%s\".\n"),
     465             :                _(PG_YESLETTER), _(PG_NOLETTER));
     466             :     }
     467             : }
     468             : 
     469             : /*
     470             :  * SetCancelConn
     471             :  *
     472             :  * Set cancelConn to point to the current database connection.
     473             :  */
     474             : void
     475        9650 : SetCancelConn(PGconn *conn)
     476             : {
     477             :     PGcancel   *oldCancelConn;
     478             : 
     479             : #ifdef WIN32
     480             :     EnterCriticalSection(&cancelConnLock);
     481             : #endif
     482             : 
     483             :     /* Free the old one if we have one */
     484        9650 :     oldCancelConn = cancelConn;
     485             : 
     486             :     /* be sure handle_sigint doesn't use pointer while freeing */
     487        9650 :     cancelConn = NULL;
     488             : 
     489        9650 :     if (oldCancelConn != NULL)
     490           0 :         PQfreeCancel(oldCancelConn);
     491             : 
     492        9650 :     cancelConn = PQgetCancel(conn);
     493             : 
     494             : #ifdef WIN32
     495             :     LeaveCriticalSection(&cancelConnLock);
     496             : #endif
     497        9650 : }
     498             : 
     499             : /*
     500             :  * ResetCancelConn
     501             :  *
     502             :  * Free the current cancel connection, if any, and set to NULL.
     503             :  */
     504             : void
     505        9650 : ResetCancelConn(void)
     506             : {
     507             :     PGcancel   *oldCancelConn;
     508             : 
     509             : #ifdef WIN32
     510             :     EnterCriticalSection(&cancelConnLock);
     511             : #endif
     512             : 
     513        9650 :     oldCancelConn = cancelConn;
     514             : 
     515             :     /* be sure handle_sigint doesn't use pointer while freeing */
     516        9650 :     cancelConn = NULL;
     517             : 
     518        9650 :     if (oldCancelConn != NULL)
     519        9650 :         PQfreeCancel(oldCancelConn);
     520             : 
     521             : #ifdef WIN32
     522             :     LeaveCriticalSection(&cancelConnLock);
     523             : #endif
     524        9650 : }
     525             : 
     526             : #ifndef WIN32
     527             : /*
     528             :  * Handle interrupt signals by canceling the current command, if a cancelConn
     529             :  * is set.
     530             :  */
     531             : static void
     532           0 : handle_sigint(SIGNAL_ARGS)
     533             : {
     534           0 :     int         save_errno = errno;
     535             :     char        errbuf[256];
     536             : 
     537             :     /* Send QueryCancel if we are processing a database query */
     538           0 :     if (cancelConn != NULL)
     539             :     {
     540           0 :         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
     541             :         {
     542           0 :             CancelRequested = true;
     543           0 :             fprintf(stderr, _("Cancel request sent\n"));
     544             :         }
     545             :         else
     546           0 :             fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
     547             :     }
     548             :     else
     549           0 :         CancelRequested = true;
     550             : 
     551           0 :     errno = save_errno;         /* just in case the write changed it */
     552           0 : }
     553             : 
     554             : void
     555         110 : setup_cancel_handler(void)
     556             : {
     557         110 :     pqsignal(SIGINT, handle_sigint);
     558         110 : }
     559             : #else                           /* WIN32 */
     560             : 
     561             : /*
     562             :  * Console control handler for Win32. Note that the control handler will
     563             :  * execute on a *different thread* than the main one, so we need to do
     564             :  * proper locking around those structures.
     565             :  */
     566             : static BOOL WINAPI
     567             : consoleHandler(DWORD dwCtrlType)
     568             : {
     569             :     char        errbuf[256];
     570             : 
     571             :     if (dwCtrlType == CTRL_C_EVENT ||
     572             :         dwCtrlType == CTRL_BREAK_EVENT)
     573             :     {
     574             :         /* Send QueryCancel if we are processing a database query */
     575             :         EnterCriticalSection(&cancelConnLock);
     576             :         if (cancelConn != NULL)
     577             :         {
     578             :             if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
     579             :             {
     580             :                 fprintf(stderr, _("Cancel request sent\n"));
     581             :                 CancelRequested = true;
     582             :             }
     583             :             else
     584             :                 fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
     585             :         }
     586             :         else
     587             :             CancelRequested = true;
     588             : 
     589             :         LeaveCriticalSection(&cancelConnLock);
     590             : 
     591             :         return TRUE;
     592             :     }
     593             :     else
     594             :         /* Return FALSE for any signals not being handled */
     595             :         return FALSE;
     596             : }
     597             : 
     598             : void
     599             : setup_cancel_handler(void)
     600             : {
     601             :     InitializeCriticalSection(&cancelConnLock);
     602             : 
     603             :     SetConsoleCtrlHandler(consoleHandler, TRUE);
     604             : }
     605             : 
     606             : #endif                          /* WIN32 */

Generated by: LCOV version 1.13