From d5a102b9a95f641a4c6be16be9f1094a4b06027c Mon Sep 17 00:00:00 2001 From: Jack Walker Date: Mon, 7 Sep 2026 21:17:09 -0400 Subject: [PATCH] Expose the PostgreSQL service name in prompts Fixes #1624. Implemented and validated with Codex. --- AUTHORS | 2 ++ changelog.rst | 2 ++ pgcli/main.py | 4 ++++ pgcli/pgclirc | 1 + tests/test_main.py | 28 ++++++++++++++++++++++++++++ 5 files changed, 37 insertions(+) diff --git a/AUTHORS b/AUTHORS index db49b9eb4..6eb4011b6 100644 --- a/AUTHORS +++ b/AUTHORS @@ -155,6 +155,8 @@ Contributors: * Chris (ChrisJr404) * Pieter Ouwerkerk (pouwerkerk) + * jackwalkerlabs (Codex-assisted) + Creator: -------- Amjith Ramanujam diff --git a/changelog.rst b/changelog.rst index fbdbc4917..feaf55580 100644 --- a/changelog.rst +++ b/changelog.rst @@ -3,6 +3,8 @@ Upcoming Features: --------- +* Add a ``\service`` prompt token for connections made using ``service=...`` + or ``PGSERVICE``. For example, set ``prompt = '\service> '``. * Add support for forcing destructive commands without confirmation. * Command line option `-y` or `--yes`. * Skips the destructive command confirmation prompt when enabled. diff --git a/pgcli/main.py b/pgcli/main.py index 8c172b85d..b07952306 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -237,6 +237,7 @@ def __init__( self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None + self.service_name = None self.watch_command = None self.force_destructive = force_destructive @@ -703,6 +704,7 @@ def connect_service(self, service, user): port=service_config.get("port"), passwd=service_config.get("password"), ) + self.service_name = service or os.getenv("PGSERVICE") def connect_uri(self, uri): kwargs = conninfo_to_dict(uri) @@ -869,6 +871,7 @@ def should_ask_for_password(exc): sys.exit(1) self.pgexecute = pgexecute + self.service_name = None def handle_editor_command(self, text): r""" @@ -1471,6 +1474,7 @@ def get_completions(self, text, cursor_positition): def get_prompt(self, string): # should be before replacing \\d string = string.replace("\\dsn_alias", self.dsn_alias or "") + string = string.replace("\\service", self.service_name or "") string = string.replace("\\t", self.now.strftime("%x %X")) string = string.replace("\\u", self.pgexecute.user or "(none)") string = string.replace("\\H", self.pgexecute.host or "(none)") diff --git a/pgcli/pgclirc b/pgcli/pgclirc index 01455c677..acd7e9960 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -188,6 +188,7 @@ verbose_errors = False # \n - Newline # \T - Transaction status: '*' if in a valid transaction, '!' if in a failed transaction, '?' if disconnected, empty otherwise # \dsn_alias - name of dsn connection string alias if -D option is used (empty otherwise) +# \service - PostgreSQL service name (empty when not connected using a service) # \x1b[...m - insert ANSI escape sequence # eg: prompt = '\x1b[35m\u@\x1b[32m\h:\x1b[36m\d>' prompt = '\u@\h:\d> ' diff --git a/tests/test_main.py b/tests/test_main.py index c8a28b419..21c1edde4 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -742,6 +742,34 @@ def test_get_prompt_with_transaction_status(transaction_indicator, expected): assert result == expected +@pytest.mark.parametrize("service_name", ["reporting", None]) +def test_get_prompt_service_name(service_name): + cli = PGCli() + cli.pgexecute = mock.MagicMock(user="user", host="localhost", short_host="localhost", dbname="db", transaction_indicator="") + cli.service_name = service_name + assert cli.get_prompt("\\service> ") == f"{service_name or ''}> " + + +@pytest.mark.parametrize("service", ["reporting", None]) +def test_connect_service_remembers_name(tmp_path, monkeypatch, service): + cli = PGCli(pgclirc_file=str(tmp_path / "rcfile")) + service_file = tmp_path / "pg_service.conf" + service_file.write_text("[reporting]\nhost=localhost\n") + monkeypatch.setenv("PGSERVICEFILE", str(service_file)) + monkeypatch.setenv("PGSERVICE", "reporting") + with mock.patch.object(cli, "connect"): + cli.connect_service(service, None) + assert cli.service_name == "reporting" + + +def test_direct_connection_clears_service_name(tmp_path): + cli = PGCli(pgclirc_file=str(tmp_path / "rcfile")) + cli.service_name = "reporting" + with mock.patch("pgcli.main.PGExecute"): + cli.connect(database="db", user="user", host="localhost") + assert cli.service_name is None + + def test_get_prompt_transaction_status_in_full_prompt(): cli = PGCli() cli.pgexecute = mock.MagicMock()