Merge branch 'main' into main

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
 
     services:
       postgres:

.github/workflows/codeql.yml

@@ -0,0 +1,41 @@
+name: "CodeQL"
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+  schedule:
+    - cron: "29 13 * * 1"
+
+jobs:
+  analyze:
+    name: Analyze
+    runs-on: ubuntu-latest
+    permissions:
+      actions: read
+      contents: read
+      security-events: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        language: [ python ]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@v2
+        with:
+          languages: ${{ matrix.language }}
+          queries: +security-and-quality
+
+      - name: Autobuild
+        uses: github/codeql-action/autobuild@v2
+
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@v2
+        with:
+          category: "/language:${{ matrix.language }}"

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.
 * Allow specifying an `alias_map_file` in the config that will use
   predetermined table aliases instead of generating aliases programmatically on
   the fly

pgcli/main.py

@@ -63,7 +63,7 @@
 )
 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 @@ def __init__(
         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:
@@ -691,7 +694,7 @@ def handle_editor_command(self, text):
             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)
@@ -718,7 +721,9 @@ def execute_command(self, text):
         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:
@@ -1041,10 +1046,17 @@ def _handle_server_closed_connection(self, text):
             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
 

setup.py

@@ -66,6 +66,7 @@
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
         "Programming Language :: SQL",
         "Topic :: Database",
         "Topic :: Database :: Front-Ends",

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py37, py38, py39, py310
+envlist = py37, py38, py39, py310, py311
 [testenv]
 deps = pytest>=2.7.0,<=3.0.7
     mock>=1.0.1