Fix: terminate running Redshift queries when SQLMesh is interrupted

sqlmesh/core/engine_adapter/redshift.py

@@ -144,6 +144,70 @@ def cursor(self) -> t.Any:
         cursor.paramstyle = "qmark"
         return cursor
 
+    def recycle(self) -> None:
+        """Closes all open connections and releases all allocated resources associated with any thread
+        except the calling one, after terminating the queries that are still running on any thread.
+
+        This includes the calling thread's query, if it has one. The calling thread can't be waiting on
+        a query while it runs this method, so a query recorded for it was probably abandoned by a signal
+        (e.g. KeyboardInterrupt). That query might still running on Redshift and holding its locks, and its
+        connection still has unread responses, so it can't be reused. The connection is replaced instead
+        of kept, which leaves the calling thread with a usable connection.
+        """
+        self._terminate_running_queries()
+        super().recycle()
+
+    def close(self) -> t.Any:
+        self._terminate_running_queries()
+        return super().close()
+
+    def _terminate_running_queries(self) -> None:
+        # Redshift keeps running a query after its client disconnects, holding its locks until it
+        # finishes, so queries that are still running (e.g. after a KeyboardInterrupt) are stopped
+        # explicitly. This must happen before the connection pool is closed, which discards the
+        # backend PIDs. PG_CANCEL_BACKEND can't cancel a query inside a transaction block, so the
+        # session is terminated instead, which also rolls back its transaction.
+        running_backend_pids = [
+            pid
+            for pid in self._connection_pool.get_all_attributes("running_query_backend_pid")
+            if pid is not None
+        ]
+        if not running_backend_pids:
+            return
+
+        if (
+            self._connection_pool.get_attribute("running_query_backend_pid") is not None
+            or self._connection_pool.is_transaction_active
+        ):
+            # The calling thread was interrupted mid-query or mid-transaction, so its connection
+            # can't be reused.
+            self._connection_pool.close()
+        for pid in running_backend_pids:
+            try:
+                self.execute(exp.select(exp.func("pg_terminate_backend", exp.Literal.number(pid))))
+                logger.info("Terminated Redshift session %s to stop its running query", pid)
+            except Exception:
+                logger.warning("Failed to terminate Redshift session %s", pid, exc_info=True)
+
+    def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None:
+        self._connection_pool.set_attribute("running_query_backend_pid", self._backend_pid())
+        try:
+            super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs)
+        except Exception:
+            self._connection_pool.set_attribute("running_query_backend_pid", None)
+            raise
+        # Not cleared on KeyboardInterrupt or SystemExit: the query is still running on the server.
+        self._connection_pool.set_attribute("running_query_backend_pid", None)
+
+    def _backend_pid(self) -> t.Optional[int]:
+        """The process ID of the current connection's Redshift session."""
+        pid = self._connection_pool.get_attribute("backend_pid")
+        if pid is None:
+            self.cursor.execute("SELECT pg_backend_pid()")
+            pid = self.cursor.fetchone()[0]
+            self._connection_pool.set_attribute("backend_pid", pid)
+        return pid
+
     def _fetch_native_df(
         self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False
     ) -> pd.DataFrame:

sqlmesh/utils/concurrency.py

@@ -1,5 +1,6 @@
 import typing as t
 from concurrent.futures import Executor, Future, ThreadPoolExecutor
+from contextlib import contextmanager
 from threading import Lock
 
 from sqlmesh.core.snapshot import SnapshotId, SnapshotInfoLike
@@ -58,7 +59,7 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]:
         if self._finished_future.done():
             self._init_state()
 
-        with ThreadPoolExecutor(max_workers=self.tasks_num) as pool:
+        with _thread_pool(self.tasks_num) as pool:
             with self._unprocessed_nodes_lock:
                 self._submit_next_nodes(pool)
             self._finished_future.result()
@@ -271,8 +272,30 @@ def _process_value(value: A, index: int) -> None:
         except Exception as ex:
             futures[index].set_exception(ex)
 
-    with ThreadPoolExecutor(max_workers=tasks_num) as pool:
+    with _thread_pool(tasks_num) as pool:
         for index, value in enumerate(values):
             pool.submit(_process_value, value, index)
 
     return [f.result() for f in futures]
+
+
+@contextmanager
+def _thread_pool(max_workers: int) -> t.Iterator[ThreadPoolExecutor]:
+    """A thread pool that, like `ThreadPoolExecutor`, waits for its tasks on exit unless the
+    calling thread is interrupted (e.g. by `KeyboardInterrupt`).
+
+    When interrupted, pending tasks are cancelled and the interrupt propagates without waiting for
+    running tasks, so that the caller can release resources such as engine adapter connections,
+    which cancels the queries those tasks are blocked on.
+    """
+    pool = ThreadPoolExecutor(max_workers=max_workers)
+    try:
+        try:
+            yield pool
+        except Exception:
+            pool.shutdown(wait=True)
+            raise
+        pool.shutdown(wait=True)
+    except BaseException:
+        pool.shutdown(wait=False, cancel_futures=True)
+        raise

tests/conftest.py

@@ -26,7 +26,11 @@
 from sqlmesh.core.config import Config, BaseDuckDBConnectionConfig, DuckDBConnectionConfig
 from sqlmesh.core.config.connection import ConnectionConfig
 from sqlmesh.core.context import Context
-from sqlmesh.core.engine_adapter import MSSQLEngineAdapter, SparkEngineAdapter
+from sqlmesh.core.engine_adapter import (
+    MSSQLEngineAdapter,
+    RedshiftEngineAdapter,
+    SparkEngineAdapter,
+)
 from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core import lineage
@@ -556,6 +560,12 @@ def _make_function(
                 "sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.catalog_support",
                 new_callable=PropertyMock(return_value=CatalogSupport.REQUIRES_SET_CATALOG),
             )
+        if isinstance(adapter, RedshiftEngineAdapter):
+            # Keeps the session lookup out of the mocked cursor's executed statements.
+            mocker.patch(
+                "sqlmesh.core.engine_adapter.redshift.RedshiftEngineAdapter._backend_pid",
+                return_value=None,
+            )
         if patch_get_data_objects:
             mocker.patch.object(adapter, "_get_data_objects", return_value=[])
         return adapter

tests/core/engine_adapter/test_redshift.py

@@ -1,4 +1,5 @@
 # type: ignore
+import threading
 import typing as t
 
 import pandas as pd  # noqa: TID253
@@ -667,6 +668,7 @@ def test_table_exists_db_table(mocker: MockerFixture):
     cursor_mock = mocker.Mock()
     connection_mock.cursor.return_value = cursor_mock
     connection_mock.cursor.return_value.fetchone.return_value = (1,)
+    mocker.patch.object(RedshiftEngineAdapter, "_backend_pid", return_value=None)
 
     adapter = RedshiftEngineAdapter(lambda: connection_mock)
     assert adapter.table_exists(table_name=exp.to_table("some_db.some_table"))
@@ -681,6 +683,7 @@ def test_table_exists_table_only(mocker: MockerFixture):
     cursor_mock = mocker.Mock()
     connection_mock.cursor.return_value = cursor_mock
     connection_mock.cursor.return_value.fetchone.return_value = None
+    mocker.patch.object(RedshiftEngineAdapter, "_backend_pid", return_value=None)
 
     adapter = RedshiftEngineAdapter(lambda: connection_mock)
     assert not adapter.table_exists(table_name=exp.to_table("some_table"))
@@ -912,3 +915,164 @@ def test_merge_logical(
         'INSERT INTO "target" ("ID", "ts") SELECT "ID", "ts" FROM (SELECT "ID" AS "ID", "ts" AS "ts", ROW_NUMBER() OVER (PARTITION BY "ID" ORDER BY "ID") AS _row_number FROM "__temp_test_abcdefgh") AS _t WHERE _row_number = 1',
         'DROP TABLE IF EXISTS "__temp_test_abcdefgh"',
     ]
+
+
+def _mock_connections(
+    mocker: MockerFixture, run_query: t.Callable[[str], None]
+) -> t.Tuple[t.Callable, t.List[t.Any]]:
+    """A connection factory whose connections have backend PIDs 100, 101, and so on."""
+    connections = []
+
+    def connect():
+        cursor = mocker.Mock()
+        cursor.fetchone.return_value = (100 + len(connections),)
+        cursor.execute.side_effect = lambda sql, **kwargs: (
+            None if sql == "SELECT pg_backend_pid()" else run_query(sql)
+        )
+        connection = mocker.NonCallableMock()
+        connection.cursor.return_value = cursor
+        connections.append(connection)
+        return connection
+
+    return connect, connections
+
+
+def _executed_sql(connection: t.Any) -> t.List[str]:
+    return [call[0][0] for call in connection.cursor.return_value.execute.call_args_list]
+
+
+def test_close_terminates_session_of_interrupted_query(mocker: MockerFixture):
+    def run_query(sql: str) -> None:
+        if sql == "SELECT long_running()":
+            raise KeyboardInterrupt
+
+    connect, connections = _mock_connections(mocker, run_query)
+    adapter = RedshiftEngineAdapter(connect)
+
+    with pytest.raises(KeyboardInterrupt):
+        adapter.execute("SELECT long_running()")
+    adapter.close()
+
+    # The interrupted connection is mid-query, so the session is terminated from a new one.
+    connections[0].close.assert_called_once()
+    assert _executed_sql(connections[1]) == [
+        "SELECT pg_backend_pid()",
+        "SELECT PG_TERMINATE_BACKEND(100)",
+    ]
+    connections[1].close.assert_called_once()
+
+
+def _recycle_then_close(adapter: RedshiftEngineAdapter) -> None:
+    # On KeyboardInterrupt, SnapshotEvaluator.concurrent_context() recycles before the context closes.
+    adapter.recycle()
+    adapter.close()
+
+
+@pytest.mark.parametrize(
+    "teardown",
+    [lambda adapter: adapter.close(), _recycle_then_close],
+    ids=["close", "recycle_then_close"],
+)
+def test_teardown_terminates_sessions_of_queries_running_in_other_threads(
+    mocker: MockerFixture, teardown: t.Callable
+):
+    query_started = threading.Event()
+    session_terminated = threading.Event()
+
+    def run_query(sql: str) -> None:
+        if sql == "SELECT long_running()":
+            query_started.set()
+            session_terminated.wait(timeout=10)
+            raise Exception("terminating connection due to administrator command")
+        if sql == "SELECT PG_TERMINATE_BACKEND(100)":
+            session_terminated.set()
+
+    connect, connections = _mock_connections(mocker, run_query)
+    adapter = RedshiftEngineAdapter(connect, multithreaded=True)
+    errors = []
+
+    def run_long_query() -> None:
+        try:
+            adapter.execute("SELECT long_running()")
+        except Exception as ex:
+            errors.append(ex)
+
+    thread = threading.Thread(target=run_long_query)
+    thread.start()
+    assert query_started.wait(timeout=10)
+    teardown(adapter)
+    thread.join(timeout=10)
+
+    assert not thread.is_alive()
+    assert len(errors) == 1
+    assert _executed_sql(connections[1]) == [
+        "SELECT pg_backend_pid()",
+        "SELECT PG_TERMINATE_BACKEND(100)",
+    ]
+    connections[0].close.assert_called_once()
+    connections[1].close.assert_called_once()
+
+
+def test_recycle_terminates_sessions_from_new_connection_after_interrupted_transaction(
+    mocker: MockerFixture,
+):
+    query_started = threading.Event()
+    session_terminated = threading.Event()
+
+    def run_query(sql: str) -> None:
+        if sql == "SELECT long_running()":
+            query_started.set()
+            session_terminated.wait(timeout=10)
+            raise Exception("terminating connection due to administrator command")
+        if sql == "SELECT PG_TERMINATE_BACKEND(100)":
+            session_terminated.set()
+
+    connect, connections = _mock_connections(mocker, run_query)
+    adapter = RedshiftEngineAdapter(connect, multithreaded=True)
+
+    def run_long_query() -> None:
+        with pytest.raises(Exception, match="terminating connection"):
+            adapter.execute("SELECT long_running()")
+
+    thread = threading.Thread(target=run_long_query)
+    thread.start()
+    assert query_started.wait(timeout=10)
+
+    with pytest.raises(KeyboardInterrupt):
+        with adapter.transaction():
+            adapter.execute("SELECT 1")
+            raise KeyboardInterrupt
+    adapter.recycle()
+    thread.join(timeout=10)
+
+    # The interrupted transaction is left open, so the session is terminated from a new connection.
+    assert not thread.is_alive()
+    connections[0].close.assert_called_once()
+    connections[1].close.assert_called_once()
+    assert _executed_sql(connections[2]) == [
+        "SELECT pg_backend_pid()",
+        "SELECT PG_TERMINATE_BACKEND(100)",
+    ]
+
+
+def test_close_without_running_queries(mocker: MockerFixture):
+    def run_query(sql: str) -> None:
+        if sql == "SELECT fail()":
+            raise Exception("query failed")
+
+    connect, connections = _mock_connections(mocker, run_query)
+    adapter = RedshiftEngineAdapter(connect)
+
+    adapter.execute("SELECT 1")
+    with pytest.raises(Exception, match="query failed"):
+        adapter.execute("SELECT fail()")
+    adapter.close()
+
+    # The backend PID is looked up once per connection, and nothing is terminated.
+    assert len(connections) == 1
+    assert _executed_sql(connections[0]) == [
+        "SELECT pg_backend_pid()",
+        "SELECT 1",
+        "SELECT fail()",
+    ]
+    connections[0].close.assert_called_once()

tests/utils/test_concurrency.py

@@ -1,12 +1,18 @@
+import signal
+import threading
+import typing as t
+
 import pytest
 from pytest_mock.plugin import MockerFixture
 
 from sqlmesh.core.snapshot import SnapshotId
 from sqlmesh.utils.concurrency import (
     NodeExecutionFailedError,
+    concurrent_apply_to_dag,
     concurrent_apply_to_snapshots,
     concurrent_apply_to_values,
 )
+from sqlmesh.utils.dag import DAG
 
 
 @pytest.mark.parametrize("tasks_num", [1, 2])
@@ -147,3 +153,48 @@ def test_concurrent_apply(tasks_num: int):
     values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
     results = concurrent_apply_to_values(values, lambda x: x * 2, tasks_num)
     assert results == [x * 2 for x in values]
+
+
+@pytest.mark.parametrize(
+    "apply",
+    [
+        lambda fn: concurrent_apply_to_values([1, 2, 3, 4], fn, 2),
+        lambda fn: concurrent_apply_to_dag(DAG({1: set(), 2: set(), 3: {1}, 4: {2}}), fn, 2),
+    ],
+    ids=["values", "dag"],
+)
+def test_interrupt_does_not_wait_for_running_tasks(apply: t.Callable):
+    started = []
+    finished = []
+    started_lock = threading.Lock()
+    both_started = threading.Event()
+    release = threading.Event()
+
+    def fn(value: int) -> None:
+        with started_lock:
+            started.append(value)
+            if len(started) == 2:
+                both_started.set()
+        release.wait(timeout=10)
+        finished.append(value)
+
+    def interrupt_main_thread() -> None:
+        both_started.wait(timeout=10)
+        signal.pthread_kill(threading.main_thread().ident, signal.SIGINT)
+
+    previous_sigint_handler = signal.signal(signal.SIGINT, signal.default_int_handler)
+    try:
+        threading.Thread(target=interrupt_main_thread).start()
+        with pytest.raises(KeyboardInterrupt):
+            apply(fn)
+        # The interrupt propagates while the running tasks are still blocked.
+        assert not finished
+    finally:
+        signal.signal(signal.SIGINT, previous_sigint_handler)
+        release.set()
+
+    for thread in threading.enumerate():
+        if thread.name.startswith("ThreadPoolExecutor"):
+            thread.join(timeout=10)
+    # Tasks that were pending when the interrupt arrived never start.
+    assert sorted(started) == [1, 2]