Recover the connection when a second Ctrl-C interrupts the query cancel

changelog.rst

@@ -6,6 +6,20 @@ Features:
 * Add a ``\service`` prompt token for connections made using ``service=...``
   or ``PGSERVICE``. For example, set ``prompt = '\service> '``.
 
+Bug fixes:
+----------
+* A second Ctrl-C while a query is being cancelled no longer leaves the
+  connection stuck. psycopg answers Ctrl-C by cancelling the query and then
+  reading the server's reply; another Ctrl-C during that cleanup, easy over an
+  SSH tunnel where it takes a round trip or two, cut it short, so the cancel
+  could be lost, the query kept running on the server and every later
+  statement failed with ``another command is already in progress``. After a
+  cancelled query pgcli now checks the connection: if it is still busy it
+  cancels again and reads the reply, and if the server does not answer within
+  5 seconds, or yet another Ctrl-C arrives, it cancels (ignoring Ctrl-C for
+  the at most 2 seconds that takes, so the query never outlives the session)
+  and closes the connection so the next statement reconnects.
+
 4.7.1 (2026-09-20)
 ==================
 

pgcli/main.py

@@ -942,6 +942,7 @@ def execute_command(self, text, handle_closed_connection=True):
             else:
                 logger.debug("cancelled query, sql: %r", text)
                 click.secho("cancelled query", err=True, fg="red")
+                self._recover_from_interrupt()
         except NotImplementedError:
             click.secho("Not Yet Implemented.", fg="yellow")
         except OperationalError as e:
@@ -1416,6 +1417,28 @@ def _evaluate_command(self, text):
 
         return output, meta_query
 
+    def _recover_from_interrupt(self):
+        """Make sure Ctrl-C did not leave the connection busy.
+
+        See PGExecute.recover_from_interrupt(). If yet another Ctrl-C lands
+        while that runs, give the connection up rather than leave it stuck.
+        """
+        if self.pgexecute is None:
+            return
+        try:
+            outcome = self.pgexecute.recover_from_interrupt()
+        except KeyboardInterrupt:
+            self.pgexecute.abandon_connection()
+            outcome = "closed"
+        if outcome == "closed":
+            self.logger.debug("Connection closed after an interrupted query")
+            click.secho(
+                "The connection was still busy after the cancel and has been closed. If the cancel did not "
+                "get through, the query may still be running on the server. The next statement reconnects.",
+                err=True,
+                fg="red",
+            )
+
     def _handle_server_closed_connection(self, text):
         """Used during CLI execution."""
         try:

pgcli/pgexecute.py

@@ -1,5 +1,9 @@
+import contextlib
 import ipaddress
 import logging
+import select
+import signal
+import time
 import traceback
 from collections import namedtuple
 import re
@@ -56,6 +60,22 @@ def register_typecasters(connection):
 
 
 # pg3: I don't know what is this
+@contextlib.contextmanager
+def _sigint_ignored():
+    """Ignore Ctrl-C inside the block (main thread only, where signals live)."""
+    try:
+        previous = signal.signal(signal.SIGINT, signal.SIG_IGN)
+    except ValueError:
+        # Not the main thread: signal handlers cannot be changed here.
+        yield
+        return
+    try:
+        yield
+    finally:
+        # None means the old handler was not set from Python.
+        signal.signal(signal.SIGINT, previous if previous is not None else signal.default_int_handler)
+
+
 class ProtocolSafeCursor(psycopg.Cursor):
     """This class wraps and suppresses Protocol Errors with pgbouncer database.
     See https://github.com/dbcli/pgcli/pull/1097.
@@ -448,6 +468,85 @@ def _must_raise(self, e):
         """
         return self.conn.closed != 0
 
+    def recover_from_interrupt(self, timeout=5.0):
+        """Leave the connection usable after Ctrl-C interrupted a query.
+
+        psycopg answers the first Ctrl-C by cancelling the query and then
+        reading the server's reply. A second Ctrl-C during that cleanup, easy
+        over an SSH tunnel where it takes a round trip or two, interrupts it:
+        the cancel may never be sent and the reply is never read, so every
+        later statement fails with "another command is already in progress"
+        while the query keeps running on the server.
+
+        Returns "idle" when the connection was not busy, "drained" once the
+        query is cancelled and its reply read, or "closed" when the server did
+        not answer within ``timeout`` and the connection had to be given up,
+        so that the next statement reconnects.
+        """
+        conn = self.conn
+        if conn is None or conn.closed or conn.pgconn.transaction_status != psycopg.pq.TransactionStatus.ACTIVE:
+            return "idle"
+        try:
+            # cancel_safe() is psycopg >= 3.2; cancel() blocks but works on 3.0.
+            if hasattr(conn, "cancel_safe"):
+                conn.cancel_safe(timeout=timeout)
+            else:
+                conn.cancel()
+        except Exception as e:
+            _logger.debug("Cancel after an interrupted query failed: %r", e)
+        if self._drain_results(timeout):
+            return "drained"
+        _logger.debug("Connection still busy %ss after the cancel, closing it", timeout)
+        conn.close()
+        return "closed"
+
+    def abandon_connection(self, timeout=2.0):
+        """Close the connection, first cancelling whatever it still runs.
+
+        For when the user interrupts the recovery itself: closing alone would
+        leave the query running on the server, which only notices the client
+        is gone when it next sends it data. Ctrl-C is ignored for the at most
+        ``timeout`` seconds the cancel takes, so it always goes out.
+        """
+        conn = self.conn
+        if conn is None or conn.closed:
+            return
+        try:
+            if conn.pgconn.transaction_status == psycopg.pq.TransactionStatus.ACTIVE:
+                with _sigint_ignored():
+                    if hasattr(conn, "cancel_safe"):
+                        conn.cancel_safe(timeout=timeout)
+                    else:
+                        conn.cancel()
+        except (Exception, KeyboardInterrupt) as e:
+            # Best effort: the connection is closed below either way.
+            _logger.debug("Cancel before closing failed: %r", e)
+        finally:
+            conn.close()
+
+    def _drain_results(self, timeout):
+        """Read and discard the reply of the query in flight.
+
+        Waits at most ``timeout`` seconds; returns True once the connection is
+        idle again.
+        """
+        pgconn = self.conn.pgconn  # type: ignore[union-attr]
+        deadline = time.monotonic() + timeout
+        try:
+            while True:
+                pgconn.consume_input()
+                if pgconn.is_busy():
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        return False
+                    select.select([pgconn.socket], [], [], remaining)
+                    continue
+                if pgconn.get_result() is None:
+                    return True
+        except psycopg.Error as e:
+            _logger.debug("Draining an interrupted query failed: %r", e)
+            return False
+
     def execute_normal_sql(self, split_sql):
         """Returns tuple (title, rows, headers, status)"""
         _logger.debug("Regular sql statement. sql: %r", split_sql)

tests/test_cancel_recovery.py

@@ -0,0 +1,117 @@
+import os
+import signal
+import time
+from unittest.mock import patch
+
+from pgcli.main import PGCli
+from utils import dbtest
+
+
+def _leave_query_unread(executor):
+    """Put the connection where a second Ctrl-C leaves it.
+
+    psycopg handles the first Ctrl-C by sending a cancel and then reading the
+    server's answer; a second Ctrl-C during that cleanup (easy over an SSH
+    tunnel, where it takes a round trip or two) interrupts it, so the query is
+    still running on the server and its answer is never read.
+    """
+    executor.conn.pgconn.send_query(b"select pg_sleep(30)")
+    raise KeyboardInterrupt
+
+
+@dbtest
+def test_double_ctrl_c_does_not_leave_the_connection_busy(executor):
+    cli = PGCli(pgexecute=executor)
+    start = time.monotonic()
+
+    with patch.object(cli, "_evaluate_command", side_effect=lambda text: _leave_query_unread(executor)):
+        cli.execute_command("select pg_sleep(30)")
+    output, query = cli._evaluate_command("select 40 + 2 as answer")
+
+    # Without the recovery every later statement failed with "another command
+    # is already in progress".
+    assert query.successful
+    assert any("42" in line for line in output)
+    # The query was cancelled on the server, not waited out.
+    assert time.monotonic() - start < 10
+
+
+@dbtest
+def test_connection_still_busy_after_the_cancel_is_closed_then_reopened(executor):
+    cli = PGCli(pgexecute=executor)
+    cli.auto_retry_closed_connection = True
+
+    with (
+        patch.object(executor, "_drain_results", return_value=False),
+        patch.object(cli, "_evaluate_command", side_effect=lambda text: _leave_query_unread(executor)),
+    ):
+        cli.execute_command("select pg_sleep(30)")
+    assert executor.conn.closed
+
+    # The next statement goes through pgcli's usual closed-connection path.
+    cli.execute_command("select 1")
+    assert not executor.conn.closed
+    output, query = cli._evaluate_command("select 40 + 2 as answer")
+    assert query.successful
+
+
+def _still_running(connection, pid):
+    """Whether backend `pid` is still running its query, given 2s to stop."""
+    for _ in range(20):
+        with connection.cursor() as cur:
+            cur.execute("select state from pg_stat_activity where pid = %s", (pid,))
+            row = cur.fetchone()
+        if row is None or row[0] != "active":
+            return False
+        time.sleep(0.1)
+    return True
+
+
+@dbtest
+def test_another_ctrl_c_during_the_recovery_closes_the_connection(executor, connection):
+    cli = PGCli(pgexecute=executor)
+    pid = executor.conn.info.backend_pid
+
+    with (
+        patch.object(executor, "recover_from_interrupt", side_effect=KeyboardInterrupt),
+        patch.object(cli, "_evaluate_command", side_effect=lambda text: _leave_query_unread(executor)),
+    ):
+        cli.execute_command("select pg_sleep(30)")
+
+    assert executor.conn.closed
+    # Closing alone would leave the sleep running on the server.
+    assert not _still_running(connection, pid)
+
+
+@dbtest
+def test_ctrl_c_during_the_last_resort_cancel_is_ignored(executor, connection):
+    # Someone hammering Ctrl-C must not cost the cancel: closing alone would
+    # leave the query running on the server.
+    pid = executor.conn.info.backend_pid
+    executor.conn.pgconn.send_query(b"select pg_sleep(30)")
+    name = "cancel_safe" if hasattr(executor.conn, "cancel_safe") else "cancel"
+    real_cancel = getattr(executor.conn, name)
+
+    def cancel_with_ctrl_c(*args, **kwargs):
+        os.kill(os.getpid(), signal.SIGINT)
+        return real_cancel(*args, **kwargs)
+
+    with patch.object(executor.conn, name, side_effect=cancel_with_ctrl_c):
+        executor.abandon_connection()
+
+    assert executor.conn.closed
+    assert not _still_running(connection, pid)
+
+
+@dbtest
+def test_idle_connection_is_left_alone(executor):
+    # Ctrl-C after the query finished (say, while printing) must not cost the
+    # session its settings.
+    executor.conn.execute("set application_name = 'kept'")
+    conn = executor.conn
+
+    assert executor.recover_from_interrupt() == "idle"
+    assert executor.conn is conn
+    with executor.conn.cursor() as cur:
+        cur.execute("show application_name")
+        assert cur.fetchone()[0] == "kept"