Add --timeout option and a connect_timeout config default

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)).

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.

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

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"))

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"