From 171f5f7feafd3531c88a4737b10e17bb50c02426 Mon Sep 17 00:00:00 2001 From: Diego Date: Wed, 19 Aug 2026 13:23:52 -0300 Subject: [PATCH] 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"