|
18 | 18 | obfuscate_process_password, |
19 | 19 | duration_in_words, |
20 | 20 | format_output, |
| 21 | + get_connect_timeout, |
21 | 22 | get_editor, |
22 | 23 | notify_callback, |
23 | 24 | PGCli, |
@@ -503,6 +504,7 @@ def test_pg_service_file(tmpdir): |
503 | 504 | "", |
504 | 505 | notify_callback, |
505 | 506 | application_name="pgcli", |
| 507 | + connect_timeout="30", |
506 | 508 | ) |
507 | 509 | del os.environ["PGPASSWORD"] |
508 | 510 | del os.environ["PGSERVICEFILE"] |
@@ -551,7 +553,7 @@ def test_application_name_db_uri(tmpdir): |
551 | 553 | mock_pgexecute.return_value = None |
552 | 554 | cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile"))) |
553 | 555 | cli.connect_uri("postgres://bar@baz.com/?application_name=cow") |
554 | | - mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow") |
| 556 | + mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow", connect_timeout="30") |
555 | 557 |
|
556 | 558 |
|
557 | 559 | @pytest.mark.parametrize( |
@@ -766,3 +768,76 @@ def test_list_databases_discards_plain_dbname(tmpdir): |
766 | 768 | path, call = _cli_conn_target(["mydb", "-l"], tmpdir) |
767 | 769 | assert path == "plain" |
768 | 770 | assert call.args[0] == "postgres" |
| 771 | + |
| 772 | + |
| 773 | +def _effective_connect_timeout(tmpdir, cli_timeout=None, dsn_timeout=None, env=None, cfgval=None): |
| 774 | + """The connect_timeout that actually reaches the connection.""" |
| 775 | + rc = str(tmpdir.join("rcfile")) |
| 776 | + with open(rc, "w") as f: |
| 777 | + f.write("[main]\n" + (f"connect_timeout = {cfgval}\n" if cfgval else "")) |
| 778 | + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} |
| 779 | + if env: |
| 780 | + environ["PGCONNECT_TIMEOUT"] = env |
| 781 | + with mock.patch.dict(os.environ, environ, clear=True): |
| 782 | + cli_obj = PGCli(pgclirc_file=rc, connect_timeout=cli_timeout) |
| 783 | + dsn = "postgresql://u@h:5432/db" + (f"?connect_timeout={dsn_timeout}" if dsn_timeout else "") |
| 784 | + captured = {} |
| 785 | + |
| 786 | + def fake(*a, **k): |
| 787 | + captured["dsn"] = k.get("dsn") or (a[5] if len(a) > 5 else None) |
| 788 | + captured["kwargs"] = k |
| 789 | + raise RuntimeError("stop") |
| 790 | + |
| 791 | + # connect() turns a failed connection into sys.exit(1); let it. |
| 792 | + with mock.patch("pgcli.main.PGExecute", side_effect=fake), pytest.raises(SystemExit): |
| 793 | + cli_obj.connect(dsn=dsn, host="h", port="5432", user="u", database="db") |
| 794 | + from_kwargs = captured.get("kwargs", {}).get("connect_timeout") |
| 795 | + return from_kwargs or conninfo_to_dict(captured.get("dsn") or "").get("connect_timeout") |
| 796 | + |
| 797 | + |
| 798 | +DSN_WITH_TIMEOUT = "postgresql://u@h:5432/db?connect_timeout=15" |
| 799 | +DSN_PLAIN = "postgresql://u@h:5432/db" |
| 800 | + |
| 801 | + |
| 802 | +@pytest.mark.parametrize( |
| 803 | + "explicit, dsn, kwargs, env, expected, why", |
| 804 | + [ |
| 805 | + (None, DSN_PLAIN, {}, None, 30, "nothing else set, so the config default applies"), |
| 806 | + (None, DSN_WITH_TIMEOUT, {}, None, None, "the connection string already says so"), |
| 807 | + (None, DSN_PLAIN, {"connect_timeout": "9"}, None, None, "the caller already says so"), |
| 808 | + (None, DSN_PLAIN, {}, "7", None, "libpq reads $PGCONNECT_TIMEOUT itself"), |
| 809 | + (None, DSN_WITH_TIMEOUT, {}, "7", None, "the connection string beats the environment"), |
| 810 | + (3, DSN_WITH_TIMEOUT, {}, "7", 3, "--timeout beats everything"), |
| 811 | + (0, DSN_WITH_TIMEOUT, {}, None, 0, "--timeout 0 is meaningful, not unset"), |
| 812 | + (None, None, {}, None, 30, "no dsn at all"), |
| 813 | + ], |
| 814 | +) |
| 815 | +def test_get_connect_timeout(explicit, dsn, kwargs, env, expected, why): |
| 816 | + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} |
| 817 | + if env: |
| 818 | + environ["PGCONNECT_TIMEOUT"] = env |
| 819 | + with mock.patch.dict(os.environ, environ, clear=True): |
| 820 | + assert get_connect_timeout(explicit, dsn, kwargs, 30) == expected, why |
| 821 | + |
| 822 | + |
| 823 | +def test_connect_timeout_config_default_reaches_the_connection(tmpdir): |
| 824 | + """The helper is actually wired into connect(): libpq's own default of 0 |
| 825 | + waits until the OS gives up, which takes minutes.""" |
| 826 | + assert _effective_connect_timeout(tmpdir) == "30" |
| 827 | + |
| 828 | + |
| 829 | +def test_connect_timeout_config_value_used(tmpdir): |
| 830 | + assert _effective_connect_timeout(tmpdir, cfgval=45) == "45" |
| 831 | + |
| 832 | + |
| 833 | +def test_connect_timeout_cli_reaches_the_connection(tmpdir): |
| 834 | + assert _effective_connect_timeout(tmpdir, cli_timeout=3, dsn_timeout=15, env="7") == "3" |
| 835 | + |
| 836 | + |
| 837 | +def test_connect_timeout_config_value_must_be_a_number(tmpdir): |
| 838 | + """A typo in the config is reported instead of being silently ignored.""" |
| 839 | + rc = str(tmpdir.join("rcfile")) |
| 840 | + with open(rc, "w") as f: |
| 841 | + f.write("[main]\nconnect_timeout = soon\n") |
| 842 | + with pytest.raises(ValueError): |
| 843 | + PGCli(pgclirc_file=rc) |
0 commit comments