feat: Add support for transaction status indicator (\T) in prompt

AUTHORS

@@ -144,6 +144,7 @@ Contributors:
     * Jay Knight (jay-knight)
     * fbdb
     * Charbel Jacquin (charbeljc)
+    * Devadathan M B (devadathanmb)
 
 Creator:
 --------

pgcli/main.py

@@ -1307,6 +1307,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("\\T", self.pgexecute.transaction_indicator)
         return string
 
     def get_last_query(self):

pgcli/pgclirc

@@ -176,6 +176,7 @@ verbose_errors = False
 # \i - Postgres PID
 # \# - "@" sign if logged in as superuser, '>' in other case
 # \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)
 # \x1b[...m - insert ANSI escape sequence
 # eg: prompt = '\x1b[35m\u@\x1b[32m\h:\x1b[36m\d>'

pgcli/pgexecute.py

@@ -298,6 +298,19 @@ def valid_transaction(self):
         status = self.conn.info.transaction_status
         return status == psycopg.pq.TransactionStatus.ACTIVE or status == psycopg.pq.TransactionStatus.INTRANS
 
+    def is_connection_closed(self):
+        return self.conn.info.transaction_status == psycopg.pq.TransactionStatus.UNKNOWN
+
+    @property
+    def transaction_indicator(self):
+        if self.is_connection_closed():
+            return "?"
+        if self.failed_transaction():
+            return "!"
+        if self.valid_transaction():
+            return "*"
+        return ""
+
     def run(
         self,
         statement,

tests/test_main.py

@@ -577,6 +577,47 @@ def test_duration_in_words(duration_in_seconds, words):
     assert duration_in_words(duration_in_seconds) == words
 
 
+@pytest.mark.parametrize(
+    "transaction_indicator,expected",
+    [
+        ("*", "*testuser"),  # valid transaction
+        ("!", "!testuser"),  # failed transaction
+        ("?", "?testuser"),  # connection closed
+        ("", "testuser"),    # idle
+    ],
+)
+def test_get_prompt_with_transaction_status(transaction_indicator, expected):
+    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
+    type(cli.pgexecute).transaction_indicator = mock.PropertyMock(return_value=transaction_indicator)
+
+    result = cli.get_prompt("\\T\\u")
+    assert result == expected
+
+
+def test_get_prompt_transaction_status_in_full_prompt():
+    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
+    type(cli.pgexecute).transaction_indicator = mock.PropertyMock(return_value="*")
+
+    result = cli.get_prompt("\\T\\u@\\h:\\d> ")
+    assert result == "*user@db.example.com:mydb> "
+
+
 @dbtest
 def test_notifications(executor):
     run(executor, "listen chan1")