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

Generated by: LCOV version 1.13