From 171f5f7feafd3531c88a4737b10e17bb50c02426 Mon Sep 17 00:00:00 2001 From: Diego Date: Wed, 19 Aug 2026 13:23:52 -0300 Subject: [PATCH 1/2] Add --timeout option and a connect_timeout config default pgcli never sets a connection timeout, so it inherits libpq's default of 0: wait until the operating system gives up on the TCP connection. Against an address that swallows SYNs, pgcli keeps hanging for minutes with no feedback. Adds a --timeout command line option and a connect_timeout config value (default 30 seconds). Precedence, highest first: 1. --timeout 2. connect_timeout in the connection string 3. $PGCONNECT_TIMEOUT 4. the config value Points 2 and 3 keep libpq's own ordering, and when only the environment variable is set nothing is injected, so libpq reads it itself. --timeout 0 is meaningful (wait forever) and is not treated as unset. The resolved value is passed as a connection parameter rather than merged into the dsn, so the user's connection string is untouched; PGExecute keeps it alongside the dsn the same way it keeps hostaddr. Adds seven tests covering each precedence combination. Two existing tests asserted the exact PGExecute call and now include the resolved timeout. --- changelog.rst | 6 +++++ pgcli/main.py | 33 ++++++++++++++++++++++++ pgcli/pgclirc | 7 ++++++ pgcli/pgexecute.py | 2 +- tests/test_main.py | 63 +++++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 109 insertions(+), 2 deletions(-) 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..9fab299e8 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -197,6 +197,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 +264,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"].get("connect_timeout", "30") self.decimal_format = c["data_formats"]["decimal"] self.float_format = c["data_formats"]["float"] self.column_date_formats = c["column_date_formats"] @@ -676,6 +680,26 @@ def connect(self, database="", host="", user="", port="", passwd="", dsn="", **k kwargs.setdefault("application_name", self.application_name) + # Connection timeout precedence, highest first: + # 1. --timeout on the command line + # 2. connect_timeout in the connection string + # 3. $PGCONNECT_TIMEOUT + # 4. the connect_timeout config value (default 30) + # libpq's own default is 0, which waits until the operating system gives + # up on the TCP connection, so an unreachable host hangs for minutes. + # The resolved value is passed as a connection parameter rather than + # merged into the dsn, leaving the user's connection string untouched. + timeout = self.connect_timeout + if timeout is None: + in_dsn = "connect_timeout" in conninfo_to_dict(dsn) if dsn else False + if not in_dsn and "connect_timeout" not in kwargs and not os.environ.get("PGCONNECT_TIMEOUT"): + try: + timeout = int(self.default_connect_timeout) + except (TypeError, ValueError): + timeout = None + 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 +1449,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 +1594,7 @@ def cli( ssh_tunnel: str, init_command: str, log_file: str, + connect_timeout: int | None, ): if version: print("Version:", __version__) @@ -1621,6 +1653,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..24451df88 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -216,7 +216,7 @@ def connect( 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")} + 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..6b3d0682c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -23,6 +23,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 +501,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 +550,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 +703,62 @@ 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") + + +def test_connect_timeout_config_default(tmpdir): + """With nothing else set, the config default is applied: 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_connection_string_wins_over_config(tmpdir): + assert _effective_connect_timeout(tmpdir, dsn_timeout=15) == "15" + + +def test_connect_timeout_connection_string_wins_over_env(tmpdir): + """libpq precedence: an explicit connect_timeout beats $PGCONNECT_TIMEOUT.""" + assert _effective_connect_timeout(tmpdir, dsn_timeout=15, env="7") == "15" + + +def test_connect_timeout_env_left_to_libpq(tmpdir): + """With only $PGCONNECT_TIMEOUT set nothing is injected, so libpq reads the + environment variable itself and the config default does not override it.""" + assert _effective_connect_timeout(tmpdir, env="7") is None + + +def test_connect_timeout_cli_overrides_everything(tmpdir): + assert _effective_connect_timeout(tmpdir, cli_timeout=3, dsn_timeout=15, env="7") == "3" + + +def test_connect_timeout_cli_zero_waits_forever(tmpdir): + """--timeout 0 is meaningful and must not be treated as unset.""" + assert _effective_connect_timeout(tmpdir, cli_timeout=0, dsn_timeout=15) == "0" From 97696ea4aa1ca4ca27b9d3bb77bad110a4b54b33 Mon Sep 17 00:00:00 2001 From: Diego Date: Fri, 21 Aug 2026 14:35:22 -0300 Subject: [PATCH 2/2] Address review: validate connect_timeout at load, extract a helper - Read the config value with as_int() in __init__ so the attribute is always an integer, and a non-numeric value is reported instead of being silently ignored. - Move the precedence rules into get_connect_timeout(), which is unit tested directly rather than through a mocked connection. - Update the stale comment in pgexecute about which parameters survive alongside a dsn. --- pgcli/main.py | 47 +++++++++++++++++++++++++------------- pgcli/pgexecute.py | 4 +++- tests/test_main.py | 57 +++++++++++++++++++++++++++++----------------- 3 files changed, 70 insertions(+), 38 deletions(-) diff --git a/pgcli/main.py b/pgcli/main.py index 9fab299e8..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 @@ -266,7 +295,7 @@ def __init__( 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"].get("connect_timeout", "30") + 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"] @@ -680,23 +709,9 @@ def connect(self, database="", host="", user="", port="", passwd="", dsn="", **k kwargs.setdefault("application_name", self.application_name) - # Connection timeout precedence, highest first: - # 1. --timeout on the command line - # 2. connect_timeout in the connection string - # 3. $PGCONNECT_TIMEOUT - # 4. the connect_timeout config value (default 30) - # libpq's own default is 0, which waits until the operating system gives - # up on the TCP connection, so an unreachable host hangs for minutes. # The resolved value is passed as a connection parameter rather than # merged into the dsn, leaving the user's connection string untouched. - timeout = self.connect_timeout - if timeout is None: - in_dsn = "connect_timeout" in conninfo_to_dict(dsn) if dsn else False - if not in_dsn and "connect_timeout" not in kwargs and not os.environ.get("PGCONNECT_TIMEOUT"): - try: - timeout = int(self.default_connect_timeout) - except (TypeError, ValueError): - timeout = None + timeout = get_connect_timeout(self.connect_timeout, dsn, kwargs, self.default_connect_timeout) if timeout is not None: kwargs["connect_timeout"] = str(timeout) diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 24451df88..46bba052d 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -215,7 +215,9 @@ def connect( new_params.update(kwargs) if new_params["dsn"]: - # When using DSN, only keep dsn, password, and hostaddr (for SSH tunnels) + # 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"]: diff --git a/tests/test_main.py b/tests/test_main.py index 6b3d0682c..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, @@ -730,35 +731,49 @@ def fake(*a, **k): return from_kwargs or conninfo_to_dict(captured.get("dsn") or "").get("connect_timeout") -def test_connect_timeout_config_default(tmpdir): - """With nothing else set, the config default is applied: 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" +DSN_WITH_TIMEOUT = "postgresql://u@h:5432/db?connect_timeout=15" +DSN_PLAIN = "postgresql://u@h:5432/db" -def test_connect_timeout_connection_string_wins_over_config(tmpdir): - assert _effective_connect_timeout(tmpdir, dsn_timeout=15) == "15" +@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_connection_string_wins_over_env(tmpdir): - """libpq precedence: an explicit connect_timeout beats $PGCONNECT_TIMEOUT.""" - assert _effective_connect_timeout(tmpdir, dsn_timeout=15, env="7") == "15" +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_env_left_to_libpq(tmpdir): - """With only $PGCONNECT_TIMEOUT set nothing is injected, so libpq reads the - environment variable itself and the config default does not override it.""" - assert _effective_connect_timeout(tmpdir, env="7") is None +def test_connect_timeout_config_value_used(tmpdir): + assert _effective_connect_timeout(tmpdir, cfgval=45) == "45" -def test_connect_timeout_cli_overrides_everything(tmpdir): +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_cli_zero_waits_forever(tmpdir): - """--timeout 0 is meaningful and must not be treated as unset.""" - assert _effective_connect_timeout(tmpdir, cli_timeout=0, dsn_timeout=15) == "0" +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)