/*------------------------------------------------------------------------- * pgut-fe.c * * Portions Copyright (c) 2008-2011, NIPPON TELEGRAPH AND TELEPHONE CORPORATION * Portions Copyright (c) 2011, Itagaki Takahiro * Portions Copyright (c) 2012-2015, The Reorg Development Team *------------------------------------------------------------------------- */ #define FRONTEND #include "pgut-fe.h" #ifdef HAVE_GETOPT_H #include #else #include #endif char *dbname = NULL; char *host = NULL; char *port = NULL; char *username = NULL; char *password = NULL; YesNo prompt_password = DEFAULT; PGconn *connection = NULL; PGconn *conn2 = NULL; worker_conns workers = { .num_workers = 0, .conns = NULL }; static bool parse_pair(const char buffer[], char key[], char value[]); static char *get_username(void); /* * Set up worker conns which will be used for concurrent index rebuilds. * 'num_workers' is the desired number of worker connections, i.e. from * --jobs flag. Due to max_connections we might not actually be able to * set up that many workers, but don't treat that as a fatal error. */ void setup_workers(int num_workers) { StringInfoData buf; int i; PGconn *conn; elog(DEBUG2, "In setup_workers(), target num_workers = %d", num_workers); if (num_workers > 1 && num_workers > workers.num_workers) { initStringInfo(&buf); if (dbname && dbname[0]) appendStringInfo(&buf, "dbname=%s ", dbname); if (host && host[0]) appendStringInfo(&buf, "host=%s ", host); if (port && port[0]) appendStringInfo(&buf, "port=%s ", port); if (username && username[0]) appendStringInfo(&buf, "user=%s ", username); if (password && password[0]) appendStringInfo(&buf, "password=%s ", password); if (workers.conns == NULL) { elog(NOTICE, "Setting up workers.conns"); workers.conns = (PGconn **) pgut_malloc(sizeof(PGconn *) * num_workers); } else { elog(ERROR, "TODO: Implement pool resizing."); } for (i = 0; i < num_workers; i++) { /* Don't prompt for password again; we should have gotten * it already from reconnect(). */ elog(DEBUG2, "Setting up worker conn %d", i); /* Don't confuse pgut_connections by using pgut_connect() * * XXX: could use PQconnectStart() and PQconnectPoll() to * open these connections in non-blocking manner. */ conn = PQconnectdb(buf.data); if (PQstatus(conn) == CONNECTION_OK) { workers.conns[i] = conn; } else { elog(WARNING, "Unable to set up worker conn #%d: %s", i, PQerrorMessage(conn)); break; } /* Hardcode a search path to avoid injections into public or pg_temp */ pgut_command(conn, "SET search_path TO pg_catalog, pg_temp, public", 0, NULL); /* Make sure each worker connection can work in non-blocking * mode. */ if (PQsetnonblocking(workers.conns[i], 1)) { elog(ERROR, "Unable to set worker connection %d " "non-blocking.", i); } } /* In case we bailed out of setting up all workers, record * how many successful worker conns we actually have. */ workers.num_workers = i; termStringInfo(&buf); } } /* Disconnect all our worker conns. */ void disconnect_workers(void) { int i; if (!(workers.num_workers)) elog(DEBUG2, "No workers to disconnect."); else { for (i = 0; i < workers.num_workers; i++) { if (workers.conns[i]) { elog(DEBUG2, "Disconnecting worker %d.", i); PQfinish(workers.conns[i]); workers.conns[i] = NULL; } else { elog(NOTICE, "Worker %d already disconnected?", i); } } workers.num_workers = 0; free(workers.conns); workers.conns = NULL; } } /* * the result is also available with the global variable 'connection'. */ void reconnect(int elevel) { StringInfoData buf; char *new_password; disconnect(); initStringInfo(&buf); if (dbname && dbname[0]) appendStringInfo(&buf, "dbname=%s ", dbname); if (host && host[0]) appendStringInfo(&buf, "host=%s ", host); if (port && port[0]) appendStringInfo(&buf, "port=%s ", port); if (username && username[0]) appendStringInfo(&buf, "user=%s ", username); if (password && password[0]) appendStringInfo(&buf, "password=%s ", password); connection = pgut_connect(buf.data, prompt_password, elevel); conn2 = pgut_connect(buf.data, prompt_password, elevel); /* update password */ if (connection) { new_password = PQpass(connection); if (new_password && new_password[0] && (password == NULL || strcmp(new_password, password) != 0)) { free(password); password = pgut_strdup(new_password); } } termStringInfo(&buf); } void disconnect(void) { if (connection) { pgut_disconnect(connection); connection = NULL; } if (conn2) { pgut_disconnect(conn2); conn2 = NULL; } disconnect_workers(); } static void option_from_env(pgut_option options[]) { size_t i; for (i = 0; options && options[i].type; i++) { pgut_option *opt = &options[i]; char name[256]; size_t j; const char *s; const char *value; if (opt->source > SOURCE_ENV || opt->allowed == SOURCE_DEFAULT || opt->allowed > SOURCE_ENV) continue; for (s = opt->lname, j = 0; *s && j < lengthof(name) - 1; s++, j++) { if (strchr("-_ ", *s)) name[j] = '_'; /* - to _ */ else name[j] = toupper(*s); } name[j] = '\0'; if ((value = getenv(name)) != NULL) pgut_setopt(opt, value, SOURCE_ENV); } } /* compare two strings ignore cases and ignore -_ */ bool pgut_keyeq(const char *lhs, const char *rhs) { for (; *lhs && *rhs; lhs++, rhs++) { if (strchr("-_ ", *lhs)) { if (!strchr("-_ ", *rhs)) return false; } else if (ToLower(*lhs) != ToLower(*rhs)) return false; } return *lhs == '\0' && *rhs == '\0'; } void pgut_setopt(pgut_option *opt, const char *optarg, pgut_optsrc src) { const char *message; if (opt == NULL) { fprintf(stderr, "Try \"%s --help\" for more information.\n", PROGRAM_NAME); exit(EINVAL); } if (opt->source > src) { /* high prior value has been set already. */ return; } else if (src >= SOURCE_CMDLINE && opt->source >= src && opt->type != 'l') { /* duplicated option in command line -- don't worry if the option * type is 'l' i.e. SimpleStringList, since we are allowed to have * multiples of these. */ message = "specified only once"; } else { /* can be overwritten if non-command line source */ opt->source = src; switch (opt->type) { case 'b': case 'B': if (optarg == NULL) { *((bool *) opt->var) = (opt->type == 'b'); return; } else if (parse_bool(optarg, (bool *) opt->var)) { return; } message = "a boolean"; break; case 'f': ((pgut_optfn) opt->var)(opt, optarg); return; case 'i': if (parse_int32(optarg, opt->var)) return; message = "a 32bit signed integer"; break; case 'l': message = "a List"; simple_string_list_append(opt->var, optarg); return; case 'u': if (parse_uint32(optarg, opt->var)) return; message = "a 32bit unsigned integer"; break; case 'I': if (parse_int64(optarg, opt->var)) return; message = "a 64bit signed integer"; break; case 'U': if (parse_uint64(optarg, opt->var)) return; message = "a 64bit unsigned integer"; break; case 's': if (opt->source != SOURCE_DEFAULT) free(*(char **) opt->var); *(char **) opt->var = pgut_strdup(optarg); return; case 't': if (parse_time(optarg, opt->var)) return; message = "a time"; break; case 'y': case 'Y': if (optarg == NULL) { *(YesNo *) opt->var = (opt->type == 'y' ? YES : NO); return; } else { bool value; if (parse_bool(optarg, &value)) { *(YesNo *) opt->var = (value ? YES : NO); return; } } message = "a boolean"; break; default: ereport(ERROR, (errcode(EINVAL), errmsg("invalid option type: %c", opt->type))); return; /* keep compiler quiet */ } } if (isprint(opt->sname)) ereport(ERROR, (errcode(EINVAL), errmsg("option -%c, --%s should be %s: '%s'", opt->sname, opt->lname, message, optarg))); else ereport(ERROR, (errcode(EINVAL), errmsg("option --%s should be %s: '%s'", opt->lname, message, optarg))); } /* * Get configuration from configuration file. */ void pgut_readopt(const char *path, pgut_option options[], int elevel) { FILE *fp; char buf[1024]; char key[1024]; char value[1024]; if (!options) return; if ((fp = pgut_fopen(path, "Rt")) == NULL) return; while (fgets(buf, lengthof(buf), fp)) { size_t i; for (i = strlen(buf); i > 0 && IsSpace(buf[i - 1]); i--) buf[i - 1] = '\0'; if (parse_pair(buf, key, value)) { for (i = 0; options[i].type; i++) { pgut_option *opt = &options[i]; if (pgut_keyeq(key, opt->lname)) { if (opt->allowed == SOURCE_DEFAULT || opt->allowed > SOURCE_FILE) elog(elevel, "option %s cannot specified in file", opt->lname); else if (opt->source <= SOURCE_FILE) pgut_setopt(opt, value, SOURCE_FILE); break; } } if (!options[i].type) elog(elevel, "invalid option \"%s\"", key); } } fclose(fp); } static const char * skip_space(const char *str, const char *line) { while (IsSpace(*str)) { str++; } return str; } static const char * get_next_token(const char *src, char *dst, const char *line) { const char *s; size_t i; size_t j; if ((s = skip_space(src, line)) == NULL) return NULL; /* parse quoted string */ if (*s == '\'') { s++; for (i = 0, j = 0; s[i] != '\0'; i++) { if (s[i] == '\\') { i++; switch (s[i]) { case 'b': dst[j] = '\b'; break; case 'f': dst[j] = '\f'; break; case 'n': dst[j] = '\n'; break; case 'r': dst[j] = '\r'; break; case 't': dst[j] = '\t'; break; case '0': case '1': case '2': case '3': case '4': case '5': case '6': case '7': { int k; long octVal = 0; for (k = 0; s[i + k] >= '0' && s[i + k] <= '7' && k < 3; k++) octVal = (octVal << 3) + (s[i + k] - '0'); i += k - 1; dst[j] = ((char) octVal); } break; default: dst[j] = s[i]; break; } } else if (s[i] == '\'') { i++; /* doubled quote becomes just one quote */ if (s[i] == '\'') dst[j] = s[i]; else break; } else dst[j] = s[i]; j++; } } else { i = j = strcspn(s, "# \n\r\t\v"); memcpy(dst, s, j); } dst[j] = '\0'; return s + i; } static bool parse_pair(const char buffer[], char key[], char value[]) { const char *start; const char *end; key[0] = value[0] = '\0'; /* * parse key */ start = buffer; if ((start = skip_space(start, buffer)) == NULL) return false; end = start + strcspn(start, "=# \n\r\t\v"); /* skip blank buffer */ if (end - start <= 0) { if (*start == '=') elog(WARNING, "syntax error in \"%s\"", buffer); return false; } /* key found */ strncpy(key, start, end - start); key[end - start] = '\0'; /* find key and value split char */ if ((start = skip_space(end, buffer)) == NULL) return false; if (*start != '=') { elog(WARNING, "syntax error in \"%s\"", buffer); return false; } start++; /* * parse value */ if ((end = get_next_token(start, value, buffer)) == NULL) return false; if ((start = skip_space(end, buffer)) == NULL) return false; if (*start != '\0' && *start != '#') { elog(WARNING, "syntax error in \"%s\"", buffer); return false; } return true; } /* * execute - Execute a SQL and return the result. */ PGresult * execute(const char *query, int nParams, const char **params) { return pgut_execute(connection, query, nParams, params); } PGresult * execute_elevel(const char *query, int nParams, const char **params, int elevel) { return pgut_execute_elevel(connection, query, nParams, params, elevel); } /* * command - Execute a SQL and discard the result. */ ExecStatusType command(const char *query, int nParams, const char **params) { return pgut_command(connection, query, nParams, params); } static void set_elevel(pgut_option *opt, const char *arg) { pgut_log_level = parse_elevel(arg); } static pgut_option default_options[] = { { 'b', 'e', "echo" , &pgut_echo }, { 'f', 'E', "elevel" , set_elevel }, { 's', 'd', "dbname" , &dbname }, { 's', 'h', "host" , &host }, { 's', 'p', "port" , &port }, { 's', 'U', "username" , &username }, { 'Y', 'w', "no-password" , &prompt_password }, { 'y', 'W', "password" , &prompt_password }, { 0 } }; static size_t option_length(const pgut_option opts[]) { size_t len; for (len = 0; opts && opts[len].type; len++) { } return len; } static pgut_option * option_find(int c, pgut_option opts1[], pgut_option opts2[]) { size_t i; for (i = 0; opts1 && opts1[i].type; i++) if (opts1[i].sname == c) return &opts1[i]; for (i = 0; opts2 && opts2[i].type; i++) if (opts2[i].sname == c) return &opts2[i]; return NULL; /* not found */ } /* * Returns the current user name. */ static char * get_username(void) { char *ret; #ifndef WIN32 struct passwd *pw; pw = getpwuid(geteuid()); ret = (pw ? pw->pw_name : NULL); #else static char username[128]; /* remains after function execute */ DWORD len = sizeof(username) - 1; if (GetUserNameA(username, &len)) ret = username; else { _dosmaperr(GetLastError()); ret = NULL; } #endif if (ret == NULL) ereport(ERROR, (errcode_errno(), errmsg("could not get current user name: "))); return ret; } static int option_has_arg(char type) { switch (type) { case 'b': case 'B': case 'y': case 'Y': return no_argument; default: return required_argument; } } static void option_copy(struct option dst[], const pgut_option opts[], size_t len) { size_t i; for (i = 0; i < len; i++) { dst[i].name = opts[i].lname; dst[i].has_arg = option_has_arg(opts[i].type); dst[i].flag = NULL; dst[i].val = opts[i].sname; } } static struct option * option_merge(const pgut_option opts1[], const pgut_option opts2[]) { struct option *result; size_t len1 = option_length(opts1); size_t len2 = option_length(opts2); size_t n = len1 + len2; result = pgut_newarray(struct option, n + 1); option_copy(result, opts1, len1); option_copy(result + len1, opts2, len2); memset(&result[n], 0, sizeof(pgut_option)); return result; } static char * longopts_to_optstring(const struct option opts[]) { size_t len; char *result; char *s; for (len = 0; opts[len].name; len++) { } result = pgut_malloc(len * 2 + 1); s = result; for (len = 0; opts[len].name; len++) { if (!isprint(opts[len].val)) continue; *s++ = opts[len].val; if (opts[len].has_arg != no_argument) *s++ = ':'; } *s = '\0'; return result; } int pgut_getopt(int argc, char **argv, pgut_option options[]) { int c; int optindex = 0; char *optstring; struct option *longopts; pgut_option *opt; pgut_init(argc, argv); /* Help message and version are handled at first. */ if (argc > 1) { if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0) { help(true); exit(0); } if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0) { printf("%s %s\n", PROGRAM_NAME, PROGRAM_VERSION); exit(0); } if (strcmp(argv[1], "--configuration") == 0) { printf("%s\n", PG_VERSION_STR); exit(0); } } /* Merge default and user options. */ longopts = option_merge(default_options, options); optstring = longopts_to_optstring(longopts); /* Assign named options */ while ((c = getopt_long(argc, argv, optstring, longopts, &optindex)) != -1) { opt = option_find(c, default_options, options); pgut_setopt(opt, optarg, SOURCE_CMDLINE); } /* Read environment variables */ option_from_env(options); (void) (dbname || (dbname = getenv("PGDATABASE")) || (dbname = getenv("PGUSER")) || (dbname = get_username())); return optind; } void help(bool details) { pgut_help(details); if (details) { printf("\nConnection options:\n"); printf(" -d, --dbname=DBNAME database to connect\n"); printf(" -h, --host=HOSTNAME database server host or socket directory\n"); printf(" -p, --port=PORT database server port\n"); printf(" -U, --username=USERNAME user name to connect as\n"); printf(" -w, --no-password never prompt for password\n"); printf(" -W, --password force password prompt\n"); } printf("\nGeneric options:\n"); if (details) { printf(" -e, --echo echo queries\n"); printf(" -E, --elevel=LEVEL set output message level\n"); } printf(" --help show this help, then exit\n"); printf(" --version output version information, then exit\n"); if (details && (PROGRAM_URL || PROGRAM_ISSUES)) { printf("\n"); if (PROGRAM_URL) printf("Read the website for details: <%s>.\n", PROGRAM_URL); if (PROGRAM_ISSUES) printf("Report bugs to <%s>.\n", PROGRAM_ISSUES); } }