Skip to content

Commit

Permalink
Use PQconnectdbParams() instead of PQconnectdb()
Browse files Browse the repository at this point in the history
  • Loading branch information
za-arthur committed Apr 11, 2024
1 parent 0d2f36c commit 135df8b
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 56 deletions.
49 changes: 26 additions & 23 deletions bin/pgut/pgut-fe.c
Original file line number Diff line number Diff line change
Expand Up @@ -44,25 +44,30 @@ static bool parse_pair(const char buffer[], char key[], char value[]);
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);
#define PARAMS_ARRAY_SIZE 6

const char *keywords[PARAMS_ARRAY_SIZE];
const char *values[PARAMS_ARRAY_SIZE];

keywords[0] = "host";
values[0] = host;
keywords[1] = "port";
values[1] = port;
keywords[2] = "user";
values[2] = username;
keywords[3] = "password";
values[3] = password;
keywords[4] = "dbname";
values[4] = dbname;
keywords[5] = NULL;
values[5] = NULL;

if (workers.conns == NULL)
{
Expand All @@ -86,11 +91,11 @@ setup_workers(int num_workers)
* 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;
}
conn = PQconnectdbParams(keywords, values, true);
if (PQstatus(conn) == CONNECTION_OK)
{
workers.conns[i] = conn;
}
else
{
elog(WARNING, "Unable to set up worker conn #%d: %s", i,
Expand All @@ -113,10 +118,8 @@ setup_workers(int num_workers)
/* 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);
}
workers.num_workers = i;
}
}

/* Disconnect all our worker conns. */
Expand Down Expand Up @@ -170,8 +173,8 @@ reconnect(int elevel)
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);
connection = pgut_connect(dbname, host, port, username, password, prompt_password, elevel);
conn2 = pgut_connect(dbname, host, port, username, password, prompt_password, elevel);

/* update password */
if (connection)
Expand Down
61 changes: 29 additions & 32 deletions bin/pgut/pgut.c
Original file line number Diff line number Diff line change
Expand Up @@ -496,34 +496,40 @@ prompt_for_password(void)


PGconn *
pgut_connect(const char *info, YesNo prompt, int elevel)
pgut_connect(const char *dbname, const char *host, const char *port,
const char *username, const char *password,
YesNo prompt, int elevel)
{
char *passwd;
StringInfoData add_pass;
char *new_password = NULL;

if (prompt == YES)
{
passwd = prompt_for_password();
initStringInfo(&add_pass);
appendStringInfoString(&add_pass, info);
appendStringInfo(&add_pass, " password='%s' ", passwd);
}
else
{
passwd = NULL;
add_pass.data = NULL;
}
new_password = prompt_for_password();

/* Start the connection. Loop until we have a password if requested by backend. */
for (;;)
{
#define PARAMS_ARRAY_SIZE 6

const char *keywords[PARAMS_ARRAY_SIZE];
const char *values[PARAMS_ARRAY_SIZE];
PGconn *conn;

CHECK_FOR_INTERRUPTS();

if (!passwd)
conn = PQconnectdb(info);
else
conn = PQconnectdb(add_pass.data);
keywords[0] = "host";
values[0] = host;
keywords[1] = "port";
values[1] = port;
keywords[2] = "user";
values[2] = username;
keywords[3] = "password";
values[3] = (new_password != NULL) ? new_password : password;
keywords[4] = "dbname";
values[4] = dbname;
keywords[5] = NULL;
values[5] = NULL;

conn = PQconnectdbParams(keywords, values, true);

if (PQstatus(conn) == CONNECTION_OK)
{
Expand All @@ -538,32 +544,23 @@ pgut_connect(const char *info, YesNo prompt, int elevel)
pgut_connections = c;
pgut_conn_unlock();

if (add_pass.data != NULL)
termStringInfo(&add_pass);
free(passwd);
free(new_password);

/* 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);

return conn;
}

if (conn && PQconnectionNeedsPassword(conn) && !passwd && prompt != NO)
if (conn && PQconnectionNeedsPassword(conn) && !new_password && prompt != NO)
{
PQfinish(conn);
passwd = prompt_for_password();
if (add_pass.data != NULL)
resetStringInfo(&add_pass);
else
initStringInfo(&add_pass);
appendStringInfoString(&add_pass, info);
appendStringInfo(&add_pass, " password='%s' ", passwd);
new_password = prompt_for_password();

continue;
}

if (add_pass.data != NULL)
termStringInfo(&add_pass);
free(passwd);
free(new_password);

ereport(elevel,
(errcode(E_PG_CONNECT),
Expand Down
4 changes: 3 additions & 1 deletion bin/pgut/pgut.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,9 @@ extern void pgut_putenv(const char *key, const char *value);
/*
* Database connections
*/
extern PGconn *pgut_connect(const char *info, YesNo prompt, int elevel);
extern PGconn *pgut_connect(const char *dbname, const char *host, const char *port,
const char *username, const char *password,
YesNo prompt, int elevel);
extern void pgut_disconnect(PGconn *conn);
extern void pgut_disconnect_all(void);
extern PGresult *pgut_execute(PGconn* conn, const char *query, int nParams, const char **params);
Expand Down

0 comments on commit 135df8b

Please sign in to comment.