LCOV - code coverage report
Current view: top level - src/bin/scripts - common.c (source / functions) Hit Total Coverage
Test: PostgreSQL 13beta1 Lines: 122 157 77.7 %
Date: 2020-05-29 01:06:25 Functions: 11 12 91.7 %
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-2020, 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/cancel.h"
      23             : #include "fe_utils/connect.h"
      24             : #include "fe_utils/string_utils.h"
      25             : 
      26             : #define ERRCODE_UNDEFINED_TABLE  "42P01"
      27             : 
      28             : /*
      29             :  * Provide strictly harmonized handling of --help and --version
      30             :  * options.
      31             :  */
      32             : void
      33         250 : handle_help_version_opts(int argc, char *argv[],
      34             :                          const char *fixed_progname, help_handler hlp)
      35             : {
      36         250 :     if (argc > 1)
      37             :     {
      38         246 :         if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
      39             :         {
      40          16 :             hlp(get_progname(argv[0]));
      41          16 :             exit(0);
      42             :         }
      43         230 :         if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
      44             :         {
      45          16 :             printf("%s (PostgreSQL) " PG_VERSION "\n", fixed_progname);
      46          16 :             exit(0);
      47             :         }
      48             :     }
      49         218 : }
      50             : 
      51             : 
      52             : /*
      53             :  * Make a database connection with the given parameters.
      54             :  *
      55             :  * An interactive password prompt is automatically issued if needed and
      56             :  * allowed by prompt_password.
      57             :  *
      58             :  * If allow_password_reuse is true, we will try to re-use any password
      59             :  * given during previous calls to this routine.  (Callers should not pass
      60             :  * allow_password_reuse=true unless reconnecting to the same database+user
      61             :  * as before, else we might create password exposure hazards.)
      62             :  */
      63             : PGconn *
      64         306 : connectDatabase(const char *dbname, const char *pghost,
      65             :                 const char *pgport, const char *pguser,
      66             :                 enum trivalue prompt_password, const char *progname,
      67             :                 bool echo, bool fail_ok, bool allow_password_reuse)
      68             : {
      69             :     PGconn     *conn;
      70             :     bool        new_pass;
      71             :     static bool have_password = false;
      72             :     static char password[100];
      73             : 
      74         306 :     if (!allow_password_reuse)
      75         176 :         have_password = false;
      76             : 
      77         306 :     if (!have_password && prompt_password == TRI_YES)
      78             :     {
      79           0 :         simple_prompt("Password: ", password, sizeof(password), false);
      80           0 :         have_password = true;
      81             :     }
      82             : 
      83             :     /*
      84             :      * Start the connection.  Loop until we have a password if requested by
      85             :      * backend.
      86             :      */
      87             :     do
      88             :     {
      89             :         const char *keywords[7];
      90             :         const char *values[7];
      91             : 
      92         306 :         keywords[0] = "host";
      93         306 :         values[0] = pghost;
      94         306 :         keywords[1] = "port";
      95         306 :         values[1] = pgport;
      96         306 :         keywords[2] = "user";
      97         306 :         values[2] = pguser;
      98         306 :         keywords[3] = "password";
      99         306 :         values[3] = have_password ? password : NULL;
     100         306 :         keywords[4] = "dbname";
     101         306 :         values[4] = dbname;
     102         306 :         keywords[5] = "fallback_application_name";
     103         306 :         values[5] = progname;
     104         306 :         keywords[6] = NULL;
     105         306 :         values[6] = NULL;
     106             : 
     107         306 :         new_pass = false;
     108         306 :         conn = PQconnectdbParams(keywords, values, true);
     109             : 
     110         306 :         if (!conn)
     111             :         {
     112           0 :             pg_log_error("could not connect to database %s: out of memory",
     113             :                          dbname);
     114           0 :             exit(1);
     115             :         }
     116             : 
     117             :         /*
     118             :          * No luck?  Trying asking (again) for a password.
     119             :          */
     120         306 :         if (PQstatus(conn) == CONNECTION_BAD &&
     121           0 :             PQconnectionNeedsPassword(conn) &&
     122             :             prompt_password != TRI_NO)
     123             :         {
     124           0 :             PQfinish(conn);
     125           0 :             simple_prompt("Password: ", password, sizeof(password), false);
     126           0 :             have_password = true;
     127           0 :             new_pass = true;
     128             :         }
     129         306 :     } while (new_pass);
     130             : 
     131             :     /* check to see that the backend connection was successfully made */
     132         306 :     if (PQstatus(conn) == CONNECTION_BAD)
     133             :     {
     134           0 :         if (fail_ok)
     135             :         {
     136           0 :             PQfinish(conn);
     137           0 :             return NULL;
     138             :         }
     139           0 :         pg_log_error("could not connect to database %s: %s",
     140             :                      dbname, PQerrorMessage(conn));
     141           0 :         exit(1);
     142             :     }
     143             : 
     144         306 :     PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL, echo));
     145             : 
     146         306 :     return conn;
     147             : }
     148             : 
     149             : /*
     150             :  * Try to connect to the appropriate maintenance database.
     151             :  */
     152             : PGconn *
     153          68 : connectMaintenanceDatabase(const char *maintenance_db,
     154             :                            const char *pghost, const char *pgport,
     155             :                            const char *pguser, enum trivalue prompt_password,
     156             :                            const char *progname, bool echo)
     157             : {
     158             :     PGconn     *conn;
     159             : 
     160             :     /* If a maintenance database name was specified, just connect to it. */
     161          68 :     if (maintenance_db)
     162           0 :         return connectDatabase(maintenance_db, pghost, pgport, pguser,
     163             :                                prompt_password, progname, echo, false, false);
     164             : 
     165             :     /* Otherwise, try postgres first and then template1. */
     166          68 :     conn = connectDatabase("postgres", pghost, pgport, pguser, prompt_password,
     167             :                            progname, echo, true, false);
     168          68 :     if (!conn)
     169           0 :         conn = connectDatabase("template1", pghost, pgport, pguser,
     170             :                                prompt_password, progname, echo, false, false);
     171             : 
     172          68 :     return conn;
     173             : }
     174             : 
     175             : /*
     176             :  * Disconnect the given connection, canceling any statement if one is active.
     177             :  */
     178             : void
     179         178 : disconnectDatabase(PGconn *conn)
     180             : {
     181             :     char        errbuf[256];
     182             : 
     183             :     Assert(conn != NULL);
     184             : 
     185         178 :     if (PQtransactionStatus(conn) == PQTRANS_ACTIVE)
     186             :     {
     187             :         PGcancel   *cancel;
     188             : 
     189           0 :         if ((cancel = PQgetCancel(conn)))
     190             :         {
     191           0 :             (void) PQcancel(cancel, errbuf, sizeof(errbuf));
     192           0 :             PQfreeCancel(cancel);
     193             :         }
     194             :     }
     195             : 
     196         178 :     PQfinish(conn);
     197         178 : }
     198             : 
     199             : /*
     200             :  * Run a query, return the results, exit program on failure.
     201             :  */
     202             : PGresult *
     203         674 : executeQuery(PGconn *conn, const char *query, bool echo)
     204             : {
     205             :     PGresult   *res;
     206             : 
     207         674 :     if (echo)
     208          94 :         printf("%s\n", query);
     209             : 
     210         674 :     res = PQexec(conn, query);
     211        1348 :     if (!res ||
     212         674 :         PQresultStatus(res) != PGRES_TUPLES_OK)
     213             :     {
     214           4 :         pg_log_error("query failed: %s", PQerrorMessage(conn));
     215           4 :         pg_log_info("query was: %s", query);
     216           4 :         PQfinish(conn);
     217           4 :         exit(1);
     218             :     }
     219             : 
     220         670 :     return res;
     221             : }
     222             : 
     223             : 
     224             : /*
     225             :  * As above for a SQL command (which returns nothing).
     226             :  */
     227             : void
     228         226 : executeCommand(PGconn *conn, const char *query, bool echo)
     229             : {
     230             :     PGresult   *res;
     231             : 
     232         226 :     if (echo)
     233          16 :         printf("%s\n", query);
     234             : 
     235         226 :     res = PQexec(conn, query);
     236         452 :     if (!res ||
     237         226 :         PQresultStatus(res) != PGRES_COMMAND_OK)
     238             :     {
     239           0 :         pg_log_error("query failed: %s", PQerrorMessage(conn));
     240           0 :         pg_log_info("query was: %s", query);
     241           0 :         PQfinish(conn);
     242           0 :         exit(1);
     243             :     }
     244             : 
     245         226 :     PQclear(res);
     246         226 : }
     247             : 
     248             : 
     249             : /*
     250             :  * As above for a SQL maintenance command (returns command success).
     251             :  * Command is executed with a cancel handler set, so Ctrl-C can
     252             :  * interrupt it.
     253             :  */
     254             : bool
     255          26 : executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
     256             : {
     257             :     PGresult   *res;
     258             :     bool        r;
     259             : 
     260          26 :     if (echo)
     261          16 :         printf("%s\n", query);
     262             : 
     263          26 :     SetCancelConn(conn);
     264          26 :     res = PQexec(conn, query);
     265          26 :     ResetCancelConn();
     266             : 
     267          26 :     r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);
     268             : 
     269          26 :     if (res)
     270          26 :         PQclear(res);
     271             : 
     272          26 :     return r;
     273             : }
     274             : 
     275             : /*
     276             :  * Consume all the results generated for the given connection until
     277             :  * nothing remains.  If at least one error is encountered, return false.
     278             :  * Note that this will block if the connection is busy.
     279             :  */
     280             : bool
     281         178 : consumeQueryResult(PGconn *conn)
     282             : {
     283         178 :     bool        ok = true;
     284             :     PGresult   *result;
     285             : 
     286         178 :     SetCancelConn(conn);
     287         358 :     while ((result = PQgetResult(conn)) != NULL)
     288             :     {
     289         180 :         if (!processQueryResult(conn, result))
     290           6 :             ok = false;
     291             :     }
     292         178 :     ResetCancelConn();
     293         178 :     return ok;
     294             : }
     295             : 
     296             : /*
     297             :  * Process (and delete) a query result.  Returns true if there's no error,
     298             :  * false otherwise -- but errors about trying to work on a missing relation
     299             :  * are reported and subsequently ignored.
     300             :  */
     301             : bool
     302        9638 : processQueryResult(PGconn *conn, PGresult *result)
     303             : {
     304             :     /*
     305             :      * If it's an error, report it.  Errors about a missing table are harmless
     306             :      * so we continue processing; but die for other errors.
     307             :      */
     308        9638 :     if (PQresultStatus(result) != PGRES_COMMAND_OK)
     309             :     {
     310           6 :         char       *sqlState = PQresultErrorField(result, PG_DIAG_SQLSTATE);
     311             : 
     312           6 :         pg_log_error("processing of database \"%s\" failed: %s",
     313             :                      PQdb(conn), PQerrorMessage(conn));
     314             : 
     315           6 :         if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) != 0)
     316             :         {
     317           6 :             PQclear(result);
     318           6 :             return false;
     319             :         }
     320             :     }
     321             : 
     322        9632 :     PQclear(result);
     323        9632 :     return true;
     324             : }
     325             : 
     326             : 
     327             : /*
     328             :  * Split TABLE[(COLUMNS)] into TABLE and [(COLUMNS)] portions.  When you
     329             :  * finish using them, pg_free(*table).  *columns is a pointer into "spec",
     330             :  * possibly to its NUL terminator.
     331             :  */
     332             : void
     333          70 : splitTableColumnsSpec(const char *spec, int encoding,
     334             :                       char **table, const char **columns)
     335             : {
     336          70 :     bool        inquotes = false;
     337          70 :     const char *cp = spec;
     338             : 
     339             :     /*
     340             :      * Find the first '(' not identifier-quoted.  Based on
     341             :      * dequote_downcase_identifier().
     342             :      */
     343         962 :     while (*cp && (*cp != '(' || inquotes))
     344             :     {
     345         892 :         if (*cp == '"')
     346             :         {
     347           6 :             if (inquotes && cp[1] == '"')
     348           2 :                 cp++;           /* pair does not affect quoting */
     349             :             else
     350           4 :                 inquotes = !inquotes;
     351           6 :             cp++;
     352             :         }
     353             :         else
     354         886 :             cp += PQmblen(cp, encoding);
     355             :     }
     356          70 :     *table = pnstrdup(spec, cp - spec);
     357          70 :     *columns = cp;
     358          70 : }
     359             : 
     360             : /*
     361             :  * Break apart TABLE[(COLUMNS)] of "spec".  With the reset_val of search_path
     362             :  * in effect, have regclassin() interpret the TABLE portion.  Append to "buf"
     363             :  * the qualified name of TABLE, followed by any (COLUMNS).  Exit on failure.
     364             :  * We use this to interpret --table=foo under the search path psql would get,
     365             :  * in advance of "ANALYZE public.foo" under the always-secure search path.
     366             :  */
     367             : void
     368          50 : appendQualifiedRelation(PQExpBuffer buf, const char *spec,
     369             :                         PGconn *conn, bool echo)
     370             : {
     371             :     char       *table;
     372             :     const char *columns;
     373             :     PQExpBufferData sql;
     374             :     PGresult   *res;
     375             :     int         ntups;
     376             : 
     377          50 :     splitTableColumnsSpec(spec, PQclientEncoding(conn), &table, &columns);
     378             : 
     379             :     /*
     380             :      * Query must remain ABSOLUTELY devoid of unqualified names.  This would
     381             :      * be unnecessary given a regclassin() variant taking a search_path
     382             :      * argument.
     383             :      */
     384          50 :     initPQExpBuffer(&sql);
     385          50 :     appendPQExpBufferStr(&sql,
     386             :                          "SELECT c.relname, ns.nspname\n"
     387             :                          " FROM pg_catalog.pg_class c,"
     388             :                          " pg_catalog.pg_namespace ns\n"
     389             :                          " WHERE c.relnamespace OPERATOR(pg_catalog.=) ns.oid\n"
     390             :                          "  AND c.oid OPERATOR(pg_catalog.=) ");
     391          50 :     appendStringLiteralConn(&sql, table, conn);
     392          50 :     appendPQExpBufferStr(&sql, "::pg_catalog.regclass;");
     393             : 
     394          50 :     executeCommand(conn, "RESET search_path;", echo);
     395             : 
     396             :     /*
     397             :      * One row is a typical result, as is a nonexistent relation ERROR.
     398             :      * regclassin() unconditionally accepts all-digits input as an OID; if no
     399             :      * relation has that OID; this query returns no rows.  Catalog corruption
     400             :      * might elicit other row counts.
     401             :      */
     402          50 :     res = executeQuery(conn, sql.data, echo);
     403          48 :     ntups = PQntuples(res);
     404          48 :     if (ntups != 1)
     405             :     {
     406           0 :         pg_log_error(ngettext("query returned %d row instead of one: %s",
     407             :                               "query returned %d rows instead of one: %s",
     408             :                               ntups),
     409             :                      ntups, sql.data);
     410           0 :         PQfinish(conn);
     411           0 :         exit(1);
     412             :     }
     413          48 :     appendPQExpBufferStr(buf,
     414          48 :                          fmtQualifiedId(PQgetvalue(res, 0, 1),
     415          48 :                                         PQgetvalue(res, 0, 0)));
     416          48 :     appendPQExpBufferStr(buf, columns);
     417          48 :     PQclear(res);
     418          48 :     termPQExpBuffer(&sql);
     419          48 :     pg_free(table);
     420             : 
     421          48 :     PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL, echo));
     422          48 : }
     423             : 
     424             : 
     425             : /*
     426             :  * Check yes/no answer in a localized way.  1=yes, 0=no, -1=neither.
     427             :  */
     428             : 
     429             : /* translator: abbreviation for "yes" */
     430             : #define PG_YESLETTER gettext_noop("y")
     431             : /* translator: abbreviation for "no" */
     432             : #define PG_NOLETTER gettext_noop("n")
     433             : 
     434             : bool
     435           0 : yesno_prompt(const char *question)
     436             : {
     437             :     char        prompt[256];
     438             : 
     439             :     /*------
     440             :        translator: This is a question followed by the translated options for
     441             :        "yes" and "no". */
     442           0 :     snprintf(prompt, sizeof(prompt), _("%s (%s/%s) "),
     443             :              _(question), _(PG_YESLETTER), _(PG_NOLETTER));
     444             : 
     445             :     for (;;)
     446           0 :     {
     447             :         char        resp[10];
     448             : 
     449           0 :         simple_prompt(prompt, resp, sizeof(resp), true);
     450             : 
     451           0 :         if (strcmp(resp, _(PG_YESLETTER)) == 0)
     452           0 :             return true;
     453           0 :         if (strcmp(resp, _(PG_NOLETTER)) == 0)
     454           0 :             return false;
     455             : 
     456           0 :         printf(_("Please answer \"%s\" or \"%s\".\n"),
     457             :                _(PG_YESLETTER), _(PG_NOLETTER));
     458             :     }
     459             : }

Generated by: LCOV version 1.13