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

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

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

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

tests/test_main.py

@@ -16,13 +16,15 @@
     obfuscate_process_password,
     duration_in_words,
     format_output,
+    get_connect_timeout,
     get_editor,
     notify_callback,
     PGCli,
     OutputSettings,
     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)