Show PostgreSQL service names in the prompt

AUTHORS

@@ -155,6 +155,8 @@ Contributors:
     * Chris (ChrisJr404)
     * Pieter Ouwerkerk (pouwerkerk)
 
+    * jackwalkerlabs (Codex-assisted)
+
 Creator:
 --------
 Amjith Ramanujam

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.

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

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> '

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