From 6609908569a61b3aff64e17bde62d8311576c12a Mon Sep 17 00:00:00 2001 From: Diego Date: Fri, 16 Oct 2026 09:37:32 -0300 Subject: [PATCH 1/2] Recover the connection when Ctrl-C interrupts the query cancel psycopg answers 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, cut it short: 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 (psycopg.pq only, at most 5s). If the server does not answer, or yet another Ctrl-C lands, it cancels (at most 2s) and closes the connection, and the next statement reconnects through the usual closed-connection path. Four tests, all failing without the change; one also fails if the last-resort path only closes the connection. --- changelog.rst | 13 +++++ pgcli/main.py | 23 +++++++++ pgcli/pgexecute.py | 78 ++++++++++++++++++++++++++++ tests/test_cancel_recovery.py | 95 +++++++++++++++++++++++++++++++++++ 4 files changed, 209 insertions(+) create mode 100644 tests/test_cancel_recovery.py diff --git a/changelog.rst b/changelog.rst index 81fbefe78..076e78d1f 100644 --- a/changelog.rst +++ b/changelog.rst @@ -6,6 +6,19 @@ 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 and closes the + connection so the next statement reconnects. + 4.7.1 (2026-09-20) ================== diff --git a/pgcli/main.py b/pgcli/main.py index 994f116e5..fdeeaad37 100644 --- a/pgcli/main.py +++ b/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: diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 1cac82eb4..484ea2056 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -1,5 +1,7 @@ import ipaddress import logging +import select +import time import traceback from collections import namedtuple import re @@ -448,6 +450,82 @@ 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. A further Ctrl-C skips the cancel. + """ + conn = self.conn + if conn is None or conn.closed: + return + try: + if conn.pgconn.transaction_status == psycopg.pq.TransactionStatus.ACTIVE: + if hasattr(conn, "cancel_safe"): + conn.cancel_safe(timeout=timeout) + else: + conn.cancel() + except (Exception, KeyboardInterrupt) as e: + _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 + 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) diff --git a/tests/test_cancel_recovery.py b/tests/test_cancel_recovery.py new file mode 100644 index 000000000..c40604d28 --- /dev/null +++ b/tests/test_cancel_recovery.py @@ -0,0 +1,95 @@ +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_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" From 709267d8736be69635a0ba327b3854ca16c7a437 Mon Sep 17 00:00:00 2001 From: Diego Date: Fri, 16 Oct 2026 14:21:43 -0300 Subject: [PATCH 2/2] Ignore Ctrl-C during the last-resort cancel When the user interrupts the recovery itself, pgcli cancels the query and closes the connection. A further Ctrl-C during that cancel skipped it, so with four quick Ctrl-C over a tunnel the query kept running on the server after the session gave up on it. Ctrl-C is now ignored for the at most 2 seconds the cancel takes (main thread only; elsewhere signal handlers cannot change). The new test fails without the change. --- changelog.rst | 5 +++-- pgcli/pgexecute.py | 33 +++++++++++++++++++++++++++------ tests/test_cancel_recovery.py | 22 ++++++++++++++++++++++ 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/changelog.rst b/changelog.rst index 076e78d1f..303d8c3d8 100644 --- a/changelog.rst +++ b/changelog.rst @@ -16,8 +16,9 @@ Bug fixes: 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 and closes the - connection so the next statement reconnects. + 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) ================== diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 484ea2056..9e2186482 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -1,6 +1,8 @@ +import contextlib import ipaddress import logging import select +import signal import time import traceback from collections import namedtuple @@ -58,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. @@ -487,18 +505,21 @@ def abandon_connection(self, timeout=2.0): 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. A further Ctrl-C skips the cancel. + 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: - if hasattr(conn, "cancel_safe"): - conn.cancel_safe(timeout=timeout) - else: - conn.cancel() + 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() @@ -509,7 +530,7 @@ def _drain_results(self, timeout): Waits at most ``timeout`` seconds; returns True once the connection is idle again. """ - pgconn = self.conn.pgconn + pgconn = self.conn.pgconn # type: ignore[union-attr] deadline = time.monotonic() + timeout try: while True: diff --git a/tests/test_cancel_recovery.py b/tests/test_cancel_recovery.py index c40604d28..00532da6b 100644 --- a/tests/test_cancel_recovery.py +++ b/tests/test_cancel_recovery.py @@ -1,3 +1,5 @@ +import os +import signal import time from unittest.mock import patch @@ -81,6 +83,26 @@ def test_another_ctrl_c_during_the_recovery_closes_the_connection(executor, conn 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