Add config option to not automatically retry connections with operational errors (#1380)

AUTHORS

@@ -125,6 +125,7 @@ Contributors:
     * Liu Zhao (astroshot)
     * Rigo Neri (rigoneri)
     * Anna Glasgall (annathyst)
+    * Andy Schoenberger (andyscho)
 
 Creator:
 --------

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):
 ===================

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

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