diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index 69fac83..7c31891 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -1516,7 +1516,7 @@ static bool do_connect(char *dbname, char *user, char *host, char *port) { PGconn *o_conn = pset.db, - *n_conn; + *n_conn = NULL; char *password = NULL; if (!dbname) @@ -1570,14 +1570,69 @@ do_connect(char *dbname, char *user, char *host, char *port) keywords[7] = NULL; values[7] = NULL; - n_conn = PQconnectdbParams(keywords, values, true); + /* attempt connection asynchronously */ + n_conn = PQconnectStartParams(keywords, values, true); + + if (sigsetjmp(sigint_interrupt_jmp, 1) != 0) + { + /* interrupted during connection attempt */ + PQfinish(n_conn); + n_conn = NULL; + } else { + int waiting = true, select_ret; + struct timeval timeout; + fd_set input_mask; + + FD_ZERO(&input_mask); + FD_SET(PQsocket(n_conn), &input_mask); + + timeout.tv_sec = 0; + timeout.tv_usec = 500000; /* 0.5 sec */ + + sigint_interrupt_enabled = true; + select_ret = select(PQsocket(n_conn) + 1, &input_mask, NULL, NULL, &timeout); + sigint_interrupt_enabled = false; + + /* don't even bother waiting, something wrong with the socket */ + if (select_ret < 0 && errno != EINTR) + { + waiting = false; + } + + while(waiting) + { + switch(PQconnectPoll(n_conn)) + { + case PGRES_POLLING_READING: + case PGRES_POLLING_WRITING: + timeout.tv_sec = 0; + timeout.tv_usec = 500000; + + sigint_interrupt_enabled = true; + select_ret = select(PQsocket(n_conn) + 1, &input_mask, NULL, NULL, &timeout); + sigint_interrupt_enabled = false; + + if (select_ret < 0 && errno != EINTR) + { + waiting = false; + } + break; + default: + waiting = false; + break; + } + } + } free(keywords); free(values); /* We can immediately discard the password -- no longer needed */ if (password) + { free(password); + password = NULL; + } if (PQstatus(n_conn) == CONNECTION_OK) break; @@ -1586,7 +1641,7 @@ do_connect(char *dbname, char *user, char *host, char *port) * Connection attempt failed; either retry the connection attempt with * a new password, or give up. */ - if (!password && PQconnectionNeedsPassword(n_conn) && pset.getPassword != TRI_NO) + if (PQconnectionNeedsPassword(n_conn) && pset.getPassword != TRI_NO) { PQfinish(n_conn); password = prompt_for_password(user); @@ -1600,7 +1655,8 @@ do_connect(char *dbname, char *user, char *host, char *port) */ if (pset.cur_cmd_interactive) { - psql_error("%s", PQerrorMessage(n_conn)); + if (n_conn) + psql_error("%s", PQerrorMessage(n_conn)); /* pset.db is left unmodified */ if (o_conn) @@ -1608,7 +1664,9 @@ do_connect(char *dbname, char *user, char *host, char *port) } else { - psql_error("\\connect: %s", PQerrorMessage(n_conn)); + if (n_conn) + psql_error("\\connect: %s", PQerrorMessage(n_conn)); + if (o_conn) { PQfinish(o_conn); diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c index 8b1864c..e53d84c 100644 --- a/src/bin/psql/startup.c +++ b/src/bin/psql/startup.c @@ -111,8 +111,6 @@ main(int argc, char *argv[]) setvbuf(stderr, NULL, _IONBF, 0); #endif - setup_cancel_handler(); - pset.progname = get_progname(argv[0]); pset.db = NULL; @@ -245,8 +243,9 @@ main(int argc, char *argv[]) } /* - * Now find something to do + * Now find something to do (and handle cancellation, if applicable) */ + setup_cancel_handler(); /* * process file given by -f