diff --git a/changelog.rst b/changelog.rst index 81fbefe78..303d8c3d8 100644 --- a/changelog.rst +++ b/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) ================== 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..9e2186482 100644 --- a/pgcli/pgexecute.py +++ b/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) diff --git a/tests/test_cancel_recovery.py b/tests/test_cancel_recovery.py new file mode 100644 index 000000000..00532da6b --- /dev/null +++ b/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"