diff --git a/AUTHORS b/AUTHORS index c6751de7..d7656eb7 100644 --- a/AUTHORS +++ b/AUTHORS @@ -125,6 +125,7 @@ Contributors: * Liu Zhao (astroshot) * Rigo Neri (rigoneri) * Anna Glasgall (annathyst) + * Andy Schoenberger (andyscho) Creator: -------- diff --git a/changelog.rst b/changelog.rst index 88b13419..6fdd2e37 100644 --- a/changelog.rst +++ b/changelog.rst @@ -8,6 +8,8 @@ Features: destructive. This would allow you to be warned on `create`, `grant`, or `insert` queries. * Destructive warnings will now include the alias dsn connection string name if provided (-D option). * pgcli.magic will now work with connection URLs that use TLS client certificates for authentication +* Have config option to retry queries on operational errors like connections being lost. + Also prevents getting stuck in a retry loop. 3.5.0 (2022/09/15): =================== diff --git a/pgcli/main.py b/pgcli/main.py index c5b5d48e..9353b6c1 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -63,7 +63,7 @@ from .config import ( ) from .key_bindings import pgcli_bindings from .packages.formatter.sqlformatter import register_new_formatter -from .packages.prompt_utils import confirm_destructive_query +from .packages.prompt_utils import confirm, confirm_destructive_query from .packages.parseutils import parse_destructive_warning from .__init__ import __version__ @@ -202,6 +202,9 @@ class PGCli: self.multiline_mode = c["main"].get("multi_line_mode", "psql") self.vi_mode = c["main"].as_bool("vi") self.auto_expand = auto_vertical_output or c["main"].as_bool("auto_expand") + self.auto_retry_closed_connection = c["main"].as_bool( + "auto_retry_closed_connection" + ) self.expanded_output = c["main"].as_bool("expand") self.pgspecial.timing_enabled = c["main"].as_bool("timing") if row_limit is not None: @@ -690,7 +693,7 @@ class PGCli: editor_command = special.editor_command(text) return text - def execute_command(self, text): + def execute_command(self, text, handle_closed_connection=True): logger = self.logger query = MetaQuery(query=text, successful=False) @@ -717,7 +720,9 @@ class PGCli: except OperationalError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) - self._handle_server_closed_connection(text) + click.secho(str(e), err=True, fg="red") + if handle_closed_connection: + self._handle_server_closed_connection(text) except (PgCliQuitError, EOFError) as e: raise except Exception as e: @@ -1040,10 +1045,17 @@ class PGCli: click.secho("Reconnecting...", fg="green") self.pgexecute.connect() click.secho("Reconnected!", fg="green") - self.execute_command(text) except OperationalError as e: click.secho("Reconnect Failed", fg="red") click.secho(str(e), err=True, fg="red") + else: + retry = self.auto_retry_closed_connection or confirm( + "Run the query from before reconnecting?" + ) + if retry: + click.secho("Running query...", fg="green") + # Don't get stuck in a retry loop + self.execute_command(text, handle_closed_connection=False) def refresh_completions(self, history=None, persist_priorities="all"): """Refresh outdated completions diff --git a/pgcli/pgclirc b/pgcli/pgclirc index 98445c14..7626db5e 100644 --- a/pgcli/pgclirc +++ b/pgcli/pgclirc @@ -35,6 +35,10 @@ expand = False # Enables auto expand mode, which is similar to `\x auto` in psql. auto_expand = False +# Auto-retry queries on connection failures and other operational errors. If +# False, will prompt to rerun the failed query instead of auto-retrying. +auto_retry_closed_connection = True + # If set to True, table suggestions will include a table alias generate_aliases = False