feat: Add transaction status support to prompt

Author: devadathanmb

pgcli/main.py

@@ -1263,6 +1263,7 @@ def get_prompt(self, string):
         string = string.replace("\\i", str(self.pgexecute.pid) or "(none)")
         string = string.replace("\\#", "#" if self.pgexecute.superuser else ">")
         string = string.replace("\\n", "\n")
+        string = string.replace("\\x", self.pgexecute.transaction_status())
         return string
 
     def get_last_query(self):

pgcli/pgclirc

@@ -184,6 +184,7 @@ verbose_errors = False
 # \i - Postgres PID
 # \# - "@" sign if logged in as superuser, '>' in other case
 # \n - Newline
+# \x - Transaction status: '*' if in a valid transaction, '!' if in a failed transaction, empty otherwise
 # \dsn_alias - name of dsn connection string alias if -D option is used (empty otherwise)
 # \x1b[...m - insert ANSI escape sequence
 # eg: prompt = '\x1b[35m\u@\x1b[32m\h:\x1b[36m\d>'

pgcli/pgexecute.py

@@ -298,6 +298,20 @@ def valid_transaction(self):
         status = self.conn.info.transaction_status
         return status == psycopg.pq.TransactionStatus.ACTIVE or status == psycopg.pq.TransactionStatus.INTRANS
 
+    def transaction_status(self):
+        """Return transaction status character for prompt.
+
+        Following psql convention:
+        - '*' when in a valid transaction (ACTIVE or INTRANS)
+        - '!' when in a failed transaction (INERROR)
+        - '' when idle
+        """
+        if self.failed_transaction():
+            return "!"
+        elif self.valid_transaction():
+            return "*"
+        return ""
+
     def run(
         self,
         statement,

tests/test_main.py

@@ -577,6 +577,52 @@ def test_duration_in_words(duration_in_seconds, words):
     assert duration_in_words(duration_in_seconds) == words
 
 
+@pytest.mark.parametrize(
+    "transaction_status,expected",
+    [
+        ("*", "*testuser"),
+        ("!", "!testuser"),
+        ("", "testuser"),
+    ],
+)
+def test_get_prompt_with_transaction_status(transaction_status, expected):
+    """Test that \\x prompt variable shows transaction status."""
+    cli = PGCli()
+    cli.pgexecute = mock.MagicMock()
+    cli.pgexecute.user = "testuser"
+    cli.pgexecute.dbname = "testdb"
+    cli.pgexecute.host = "localhost"
+    cli.pgexecute.short_host = "localhost"
+    cli.pgexecute.port = 5432
+    cli.pgexecute.pid = 12345
+    cli.pgexecute.superuser = False
+
+    with mock.patch.object(
+        cli.pgexecute, "transaction_status", return_value=transaction_status
+    ):
+        result = cli.get_prompt("\\x\\u")
+    assert result == expected
+
+
+def test_get_prompt_transaction_status_in_full_prompt():
+    """Test that \\x works in a full prompt format."""
+    cli = PGCli()
+    cli.pgexecute = mock.MagicMock()
+    cli.pgexecute.user = "user"
+    cli.pgexecute.dbname = "mydb"
+    cli.pgexecute.host = "db.example.com"
+    cli.pgexecute.short_host = "db.example.com"
+    cli.pgexecute.port = 5432
+    cli.pgexecute.pid = 12345
+    cli.pgexecute.superuser = False
+
+    with mock.patch.object(
+        cli.pgexecute, "transaction_status", return_value="*"
+    ):
+        result = cli.get_prompt("\\x\\u@\\h:\\d> ")
+    assert result == "*user@db.example.com:mydb> "
+
+
 @dbtest
 def test_notifications(executor):
     run(executor, "listen chan1")