|
16 | 16 | obfuscate_process_password, |
17 | 17 | duration_in_words, |
18 | 18 | format_output, |
| 19 | + get_connect_timeout, |
19 | 20 | get_editor, |
20 | 21 | notify_callback, |
21 | 22 | PGCli, |
22 | 23 | OutputSettings, |
23 | 24 | COLOR_CODE_REGEX, |
24 | 25 | ) |
25 | 26 | from pgcli.pgexecute import PGExecute |
| 27 | +from psycopg.conninfo import conninfo_to_dict |
26 | 28 | from pgspecial.main import PAGER_OFF, PAGER_LONG_OUTPUT, PAGER_ALWAYS |
27 | 29 | from utils import dbtest, run |
28 | 30 | from collections import namedtuple |
@@ -500,6 +502,7 @@ def test_pg_service_file(tmpdir): |
500 | 502 | "", |
501 | 503 | notify_callback, |
502 | 504 | application_name="pgcli", |
| 505 | + connect_timeout="30", |
503 | 506 | ) |
504 | 507 | del os.environ["PGPASSWORD"] |
505 | 508 | del os.environ["PGSERVICEFILE"] |
@@ -548,7 +551,7 @@ def test_application_name_db_uri(tmpdir): |
548 | 551 | mock_pgexecute.return_value = None |
549 | 552 | cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile"))) |
550 | 553 | cli.connect_uri("postgres://bar@baz.com/?application_name=cow") |
551 | | - mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow") |
| 554 | + mock_pgexecute.assert_called_with("bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow", connect_timeout="30") |
552 | 555 |
|
553 | 556 |
|
554 | 557 | @pytest.mark.parametrize( |
@@ -701,3 +704,76 @@ def test_get_editor_precedence(): |
701 | 704 | # Nothing set -> None, so click uses its platform default. |
702 | 705 | with mock.patch.dict(os.environ, {}, clear=True): |
703 | 706 | assert get_editor() is None |
| 707 | + |
| 708 | + |
| 709 | +def _effective_connect_timeout(tmpdir, cli_timeout=None, dsn_timeout=None, env=None, cfgval=None): |
| 710 | + """The connect_timeout that actually reaches the connection.""" |
| 711 | + rc = str(tmpdir.join("rcfile")) |
| 712 | + with open(rc, "w") as f: |
| 713 | + f.write("[main]\n" + (f"connect_timeout = {cfgval}\n" if cfgval else "")) |
| 714 | + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} |
| 715 | + if env: |
| 716 | + environ["PGCONNECT_TIMEOUT"] = env |
| 717 | + with mock.patch.dict(os.environ, environ, clear=True): |
| 718 | + cli_obj = PGCli(pgclirc_file=rc, connect_timeout=cli_timeout) |
| 719 | + dsn = "postgresql://u@h:5432/db" + (f"?connect_timeout={dsn_timeout}" if dsn_timeout else "") |
| 720 | + captured = {} |
| 721 | + |
| 722 | + def fake(*a, **k): |
| 723 | + captured["dsn"] = k.get("dsn") or (a[5] if len(a) > 5 else None) |
| 724 | + captured["kwargs"] = k |
| 725 | + raise RuntimeError("stop") |
| 726 | + |
| 727 | + # connect() turns a failed connection into sys.exit(1); let it. |
| 728 | + with mock.patch("pgcli.main.PGExecute", side_effect=fake), pytest.raises(SystemExit): |
| 729 | + cli_obj.connect(dsn=dsn, host="h", port="5432", user="u", database="db") |
| 730 | + from_kwargs = captured.get("kwargs", {}).get("connect_timeout") |
| 731 | + return from_kwargs or conninfo_to_dict(captured.get("dsn") or "").get("connect_timeout") |
| 732 | + |
| 733 | + |
| 734 | +DSN_WITH_TIMEOUT = "postgresql://u@h:5432/db?connect_timeout=15" |
| 735 | +DSN_PLAIN = "postgresql://u@h:5432/db" |
| 736 | + |
| 737 | + |
| 738 | +@pytest.mark.parametrize( |
| 739 | + "explicit, dsn, kwargs, env, expected, why", |
| 740 | + [ |
| 741 | + (None, DSN_PLAIN, {}, None, 30, "nothing else set, so the config default applies"), |
| 742 | + (None, DSN_WITH_TIMEOUT, {}, None, None, "the connection string already says so"), |
| 743 | + (None, DSN_PLAIN, {"connect_timeout": "9"}, None, None, "the caller already says so"), |
| 744 | + (None, DSN_PLAIN, {}, "7", None, "libpq reads $PGCONNECT_TIMEOUT itself"), |
| 745 | + (None, DSN_WITH_TIMEOUT, {}, "7", None, "the connection string beats the environment"), |
| 746 | + (3, DSN_WITH_TIMEOUT, {}, "7", 3, "--timeout beats everything"), |
| 747 | + (0, DSN_WITH_TIMEOUT, {}, None, 0, "--timeout 0 is meaningful, not unset"), |
| 748 | + (None, None, {}, None, 30, "no dsn at all"), |
| 749 | + ], |
| 750 | +) |
| 751 | +def test_get_connect_timeout(explicit, dsn, kwargs, env, expected, why): |
| 752 | + environ = {k: v for k, v in os.environ.items() if k != "PGCONNECT_TIMEOUT"} |
| 753 | + if env: |
| 754 | + environ["PGCONNECT_TIMEOUT"] = env |
| 755 | + with mock.patch.dict(os.environ, environ, clear=True): |
| 756 | + assert get_connect_timeout(explicit, dsn, kwargs, 30) == expected, why |
| 757 | + |
| 758 | + |
| 759 | +def test_connect_timeout_config_default_reaches_the_connection(tmpdir): |
| 760 | + """The helper is actually wired into connect(): libpq's own default of 0 |
| 761 | + waits until the OS gives up, which takes minutes.""" |
| 762 | + assert _effective_connect_timeout(tmpdir) == "30" |
| 763 | + |
| 764 | + |
| 765 | +def test_connect_timeout_config_value_used(tmpdir): |
| 766 | + assert _effective_connect_timeout(tmpdir, cfgval=45) == "45" |
| 767 | + |
| 768 | + |
| 769 | +def test_connect_timeout_cli_reaches_the_connection(tmpdir): |
| 770 | + assert _effective_connect_timeout(tmpdir, cli_timeout=3, dsn_timeout=15, env="7") == "3" |
| 771 | + |
| 772 | + |
| 773 | +def test_connect_timeout_config_value_must_be_a_number(tmpdir): |
| 774 | + """A typo in the config is reported instead of being silently ignored.""" |
| 775 | + rc = str(tmpdir.join("rcfile")) |
| 776 | + with open(rc, "w") as f: |
| 777 | + f.write("[main]\nconnect_timeout = soon\n") |
| 778 | + with pytest.raises(ValueError): |
| 779 | + PGCli(pgclirc_file=rc) |
0 commit comments