diff --git a/changelog.rst b/changelog.rst index 6e4c47c36..364257fbb 100644 --- a/changelog.rst +++ b/changelog.rst @@ -10,6 +10,12 @@ Bug fixes: Features: --------- +* Add a ``--timeout`` command line option and a ``connect_timeout`` config value + (default 30 seconds) for the connection timeout. Precedence, highest first: + ``--timeout``, then a ``connect_timeout`` in the connection string, then + ``$PGCONNECT_TIMEOUT``, then the config value. libpq's own default is 0, + which waits until the operating system gives up on the TCP connection, so an + unreachable host used to hang for minutes. * Honor the ``PSQL_EDITOR`` environment variable when opening the external editor (``\\e``, ``\\ev``, ``\\ef``, ``\\ne``), matching psql's precedence of ``PSQL_EDITOR``, then ``EDITOR``, then ``VISUAL`` ([issue 1398](https://github.com/dbcli/pgcli/issues/1398)). diff --git a/pgcli/main.py b/pgcli/main.py index e3ba5bfa8..433c25357 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -157,6 +157,35 @@ def get_editor(): return os.environ.get("PSQL_EDITOR") or os.environ.get("EDITOR") or os.environ.get("VISUAL") or None +def get_connect_timeout(explicit, dsn, kwargs, default): + """Pick the connection timeout to apply, in seconds. + + Precedence, highest first: + + 1. ``explicit``, i.e. ``--timeout`` on the command line + 2. ``connect_timeout`` in the connection string, or in ``kwargs`` + 3. ``$PGCONNECT_TIMEOUT`` + 4. ``default``, the ``connect_timeout`` config value + + Returns ``None`` when the user already stated a timeout by one of the + means we must not override, in which case the caller leaves the + connection parameters alone and libpq reads it from where it already is. + + A default matters because libpq's own is 0, which waits until the + operating system gives up on the TCP connection, so an unreachable host + hangs for minutes. + """ + if explicit is not None: + return explicit + if "connect_timeout" in kwargs: + return None + if dsn and "connect_timeout" in conninfo_to_dict(dsn): + return None + if os.environ.get("PGCONNECT_TIMEOUT"): + return None + return default + + class PGCli: default_prompt = "\\u@\\h:\\d> " max_len_prompt = 30 @@ -197,6 +226,7 @@ def __init__( auto_vertical_output=False, warn=None, ssh_tunnel_url: str | None = None, + connect_timeout: int | None = None, log_file: str | None = None, ): self.force_passwd_prompt = force_passwd_prompt @@ -263,6 +293,9 @@ def __init__( self.prompt_format = prompt if prompt is not None else c["main"].get("prompt", self.default_prompt) self.prompt_dsn_format = prompt_dsn self.on_error = c["main"]["on_error"].upper() + # Connection timeout, in seconds. See connect() for the precedence. + self.connect_timeout = connect_timeout + self.default_connect_timeout = c["main"].as_int("connect_timeout") self.decimal_format = c["data_formats"]["decimal"] self.float_format = c["data_formats"]["float"] self.column_date_formats = c["column_date_formats"] @@ -676,6 +709,12 @@ def connect(self, database="", host="", user="", port="", passwd="", dsn="", **k kwargs.setdefault("application_name", self.application_name) + # The resolved value is passed as a connection parameter rather than + # merged into the dsn, leaving the user's connection string untouched. + timeout = get_connect_timeout(self.connect_timeout, dsn, kwargs, self.default_connect_timeout) + if timeout is not None: + kwargs["connect_timeout"] = str(timeout) + # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: @@ -1425,6 +1464,13 @@ def echo_via_pager(self, text, color=None): help="Username to connect to the postgres database.", ) @click.option("-u", "--user", "username_opt", help="Username to connect to the postgres database.") +@click.option( + "--timeout", + "connect_timeout", + type=click.INT, + default=None, + help="Seconds to wait for a connection before giving up (0 waits forever). Overrides the connection string and $PGCONNECT_TIMEOUT.", +) @click.option( "-W", "--password", @@ -1563,6 +1609,7 @@ def cli( ssh_tunnel: str, init_command: str, log_file: str, + connect_timeout: int | None, ): if version: print("Version:", __version__) @@ -1621,6 +1668,7 @@ def cli( warn=warn, ssh_tunnel_url=ssh_tunnel, log_file=log_file, + connect_timeout=connect_timeout, ) # Choose which ever one has a valid value. diff --git a/pgcli/pgclirc b/pgcli/pgclirc index 569705564..01455c677 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -146,6 +146,13 @@ syntax_style = default # for end are available in the REPL. vi = False +# Seconds to wait for a connection before giving up. Only applies when nothing +# else specifies a timeout: an explicit --timeout on the command line wins, then +# a connect_timeout in the connection string, then $PGCONNECT_TIMEOUT. Use 0 to +# wait forever, which is libpq's own default (the OS then gives up on the TCP +# connection after a few minutes). +connect_timeout = 30 + # Error handling # When one of multiple SQL statements causes an error, choose to either # continue executing the remaining statements, or stopping diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 578f8291d..46bba052d 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -215,8 +215,10 @@ def connect( new_params.update(kwargs) if new_params["dsn"]: - # When using DSN, only keep dsn, password, and hostaddr (for SSH tunnels) - new_params = {k: v for k, v in new_params.items() if k in ("dsn", "password", "hostaddr")} + # When using a DSN, the connection details all live in the dsn + # itself. Only keep the parameters that have to stay outside it: + # the password, hostaddr (for SSH tunnels) and connect_timeout. + new_params = {k: v for k, v in new_params.items() if k in ("dsn", "password", "hostaddr", "connect_timeout")} if new_params["password"]: new_params["dsn"] = make_conninfo(new_params["dsn"], password=new_params.pop("password")) diff --git a/tests/test_main.py b/tests/test_main.py index ba990631a..78dbeeeb7 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -16,6 +16,7 @@ obfuscate_process_password, duration_in_words, format_output, + get_connect_timeout, get_editor, notify_callback, PGCli, @@ -23,6 +24,7 @@ COLOR_CODE_REGEX, ) from pgcli.pgexecute import PGExecute +from psycopg.conninfo import conninfo_to_dict from pgspecial.main import PAGER_OFF, PAGER_LONG_OUTPUT, PAGER_ALWAYS from utils import dbtest, run from collections import namedtuple @@ -500,6 +502,7 @@ def test_pg_service_file(tmpdir): "", notify_callback, application_name="pgcli", + connect_timeout="30", ) del os.environ["PGPASSWORD"] del os.environ["PGSERVICEFILE"] @@ -548,7 +551,7 @@ def test_application_name_db_uri(tmpdir): mock_pgexecute.return_value = None cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile"))) cli.connect_uri("postgres://bar@baz.com/?application_name=cow") - mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow") + mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow", connect_timeout="30") @pytest.mark.parametrize( @@ -701,3 +704,76 @@ def test_get_editor_precedence(): # Nothing set -> None, so click uses its platform default. with mock.patch.dict(os.environ, {}, clear=True): assert get_editor() is None + + +def _effective_connect_timeout(tmpdir, cli_timeout=None, dsn_timeout=None, env=None, cfgval=None): + """The connect_timeout that actually reaches the connection.""" + rc = str(tmpdir.join("rcfile")) + with open(rc, "w") as f: + f.write("[main]\n" + (f"connect_timeout = {cfgval}\n" if cfgval else "")) + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} + if env: + environ["PGCONNECT_TIMEOUT"] = env + with mock.patch.dict(os.environ, environ, clear=True): + cli_obj = PGCli(pgclirc_file=rc, connect_timeout=cli_timeout) + dsn = "postgresql://u@h:5432/db" + (f"?connect_timeout={dsn_timeout}" if dsn_timeout else "") + captured = {} + + def fake(*a, **k): + captured["dsn"] = k.get("dsn") or (a[5] if len(a) > 5 else None) + captured["kwargs"] = k + raise RuntimeError("stop") + + # connect() turns a failed connection into sys.exit(1); let it. + with mock.patch("pgcli.main.PGExecute", side_effect=fake), pytest.raises(SystemExit): + cli_obj.connect(dsn=dsn, host="h", port="5432", user="u", database="db") + from_kwargs = captured.get("kwargs", {}).get("connect_timeout") + return from_kwargs or conninfo_to_dict(captured.get("dsn") or "").get("connect_timeout") + + +DSN_WITH_TIMEOUT = "postgresql://u@h:5432/db?connect_timeout=15" +DSN_PLAIN = "postgresql://u@h:5432/db" + + +@pytest.mark.parametrize( + "explicit, dsn, kwargs, env, expected, why", + [ + (None, DSN_PLAIN, {}, None, 30, "nothing else set, so the config default applies"), + (None, DSN_WITH_TIMEOUT, {}, None, None, "the connection string already says so"), + (None, DSN_PLAIN, {"connect_timeout": "9"}, None, None, "the caller already says so"), + (None, DSN_PLAIN, {}, "7", None, "libpq reads $PGCONNECT_TIMEOUT itself"), + (None, DSN_WITH_TIMEOUT, {}, "7", None, "the connection string beats the environment"), + (3, DSN_WITH_TIMEOUT, {}, "7", 3, "--timeout beats everything"), + (0, DSN_WITH_TIMEOUT, {}, None, 0, "--timeout 0 is meaningful, not unset"), + (None, None, {}, None, 30, "no dsn at all"), + ], +) +def test_get_connect_timeout(explicit, dsn, kwargs, env, expected, why): + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} + if env: + environ["PGCONNECT_TIMEOUT"] = env + with mock.patch.dict(os.environ, environ, clear=True): + assert get_connect_timeout(explicit, dsn, kwargs, 30) == expected, why + + +def test_connect_timeout_config_default_reaches_the_connection(tmpdir): + """The helper is actually wired into connect(): libpq's own default of 0 + waits until the OS gives up, which takes minutes.""" + assert _effective_connect_timeout(tmpdir) == "30" + + +def test_connect_timeout_config_value_used(tmpdir): + assert _effective_connect_timeout(tmpdir, cfgval=45) == "45" + + +def test_connect_timeout_cli_reaches_the_connection(tmpdir): + assert _effective_connect_timeout(tmpdir, cli_timeout=3, dsn_timeout=15, env="7") == "3" + + +def test_connect_timeout_config_value_must_be_a_number(tmpdir): + """A typo in the config is reported instead of being silently ignored.""" + rc = str(tmpdir.join("rcfile")) + with open(rc, "w") as f: + f.write("[main]\nconnect_timeout = soon\n") + with pytest.raises(ValueError): + PGCli(pgclirc_file=rc)