diff --git a/AUTHORS b/AUTHORS index 771de13f1..2eb2ea70b 100644 --- a/AUTHORS +++ b/AUTHORS @@ -144,6 +144,7 @@ Contributors: * Jay Knight (jay-knight) * fbdb * Charbel Jacquin (charbeljc) + * Devadathan M B (devadathanmb) Creator: -------- diff --git a/changelog.rst b/changelog.rst index 45ee4adfa..fdfcde538 100644 --- a/changelog.rst +++ b/changelog.rst @@ -1,3 +1,10 @@ +Upcoming (TBD) +============== + +Features: +--------- +* Add support for `\\T` prompt escape sequence to display transaction status (similar to psql's `%x`). + 4.4.0 (2025-12-24) ================== diff --git a/pgcli/main.py b/pgcli/main.py index de09cc1af..913228b33 100644 --- a/pgcli/main.py +++ b/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): diff --git a/pgcli/pgclirc b/pgcli/pgclirc index cac738937..35ff41c5a 100644 --- a/pgcli/pgclirc +++ b/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>' diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 9918da330..2864c8645 100644 --- a/pgcli/pgexecute.py +++ b/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, diff --git a/tests/test_main.py b/tests/test_main.py index 5cf1d09f8..52415a008 100644 --- a/tests/test_main.py +++ b/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 + cli.pgexecute.transaction_indicator = 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 + cli.pgexecute.transaction_indicator = "*" + + result = cli.get_prompt("\\T\\u@\\h:\\d> ") + assert result == "*user@db.example.com:mydb> " + + @dbtest def test_notifications(executor): run(executor, "listen chan1")