diff options
Diffstat (limited to '')
| -rw-r--r-- | pgcli/pgexecute.py | 17 |
1 files changed, 16 insertions, 1 deletions
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 497d681..e091757 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -1,3 +1,4 @@ +import ipaddress import logging import traceback from collections import namedtuple @@ -166,6 +167,7 @@ class PGExecute: host=None, port=None, dsn=None, + notify_callback=None, **kwargs, ): self._conn_params = {} @@ -178,6 +180,7 @@ class PGExecute: self.port = None self.server_version = None self.extra_args = None + self.notify_callback = notify_callback self.connect(database, user, password, host, port, dsn, **kwargs) self.reset_expanded = None @@ -236,6 +239,9 @@ class PGExecute: self.conn = conn self.conn.autocommit = True + if self.notify_callback is not None: + self.conn.add_notify_handler(self.notify_callback) + # When we connect using a DSN, we don't really know what db, # user, etc. we connected to. Let's read it. # Note: moved this after setting autocommit because of #664. @@ -273,6 +279,11 @@ class PGExecute: @property def short_host(self): + try: + ipaddress.ip_address(self.host) + return self.host + except ValueError: + pass if "," in self.host: host, _, _ = self.host.partition(",") else: @@ -431,7 +442,11 @@ class PGExecute: def handle_notices(n): nonlocal title - title = f"{n.message_primary}\n{n.message_detail}\n{title}" + title = f"{title}" + if n.message_primary is not None: + title = f"{title}\n{n.message_primary}" + if n.message_detail is not None: + title = f"{title}\n{n.message_detail}" self.conn.add_notice_handler(handle_notices) |