From 7e4f720e62a03e913e171cdce8c42b21ad4a82f6 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Tue, 13 Oct 2026 17:49:29 -0700 Subject: [PATCH 01/10] Fix: terminate running Redshift queries when SQLMesh is interrupted --- sqlmesh/cli/main.py | 8 ++ sqlmesh/core/engine_adapter/redshift.py | 44 +++++++++ sqlmesh/utils/concurrency.py | 27 ++++- tests/cli/test_cli.py | 20 ++++ tests/conftest.py | 12 ++- tests/core/engine_adapter/test_redshift.py | 109 +++++++++++++++++++++ tests/utils/test_concurrency.py | 51 ++++++++++ 7 files changed, 268 insertions(+), 3 deletions(-) diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index fc6cad8011..b542fd7418 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -2,7 +2,9 @@ import logging import os +import signal import sys +import threading import typing as t from pathlib import Path @@ -120,6 +122,12 @@ def cli( if "--help" in sys.argv: return + if threading.current_thread() is threading.main_thread(): + # Handle SIGTERM (e.g. from CI cancellation or container shutdown) like Ctrl-C, so that the + # context is closed and running queries are cancelled before exiting. + previous_sigterm_handler = signal.signal(signal.SIGTERM, signal.default_int_handler) + ctx.call_on_close(lambda: signal.signal(signal.SIGTERM, previous_sigterm_handler)) + configure_logging( debug, log_to_stdout, diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py index 8f2f098d08..80c1a886d6 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -144,6 +144,50 @@ def cursor(self) -> t.Any: cursor.paramstyle = "qmark" return cursor + def close(self) -> t.Any: + # 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. 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 running_backend_pids: + if self._connection_pool.get_attribute("running_query_backend_pid") is not None: + # The calling thread was interrupted mid-query, 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) + + return super().close() + + 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: diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index c5f78645f6..16adc8122b 100644 --- a/sqlmesh/utils/concurrency.py +++ b/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 diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index f1540727b1..36e69ca4a5 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -1,6 +1,7 @@ import json import os import pytest +import signal import string import time_machine from os import getcwd, path, remove @@ -2976,3 +2977,22 @@ def execute(context: ExecutionContext, **kwargs: t.Any) -> pd.DataFrame: assert result.exit_code != 0 assert "is not a known model or test file" in result.output mock.assert_not_called() + + +def test_sigterm_handled_like_keyboard_interrupt(runner, tmp_path, mocker): + create_example_project(tmp_path) + sigterm_handlers = [] + mocker.patch.object( + Context, + "plan", + side_effect=lambda *args, **kwargs: sigterm_handlers.append( + signal.getsignal(signal.SIGTERM) + ), + ) + previous_handler = signal.getsignal(signal.SIGTERM) + + result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "plan"]) + + assert result.exit_code == 0, result.output + assert sigterm_handlers == [signal.default_int_handler] + assert signal.getsignal(signal.SIGTERM) is previous_handler diff --git a/tests/conftest.py b/tests/conftest.py index 1f7f6d9325..499f19c5c2 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index 7211244e91..fcd8406f89 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/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,109 @@ 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 test_close_terminates_sessions_of_queries_running_in_other_threads(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) + 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) + adapter.close() + 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_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() diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 5e1e4326f7..1db83c3d1d 100644 --- a/tests/utils/test_concurrency.py +++ b/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] From 438bf6bfed9785e30168b233a5362eeaa125c1ac Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Thu, 15 Oct 2026 10:55:51 -0700 Subject: [PATCH 02/10] Fix: terminate running Redshift queries on recycle as well as close On Ctrl-C, SnapshotEvaluator.concurrent_context() recycles the adapter before the context closes it. ThreadLocalConnectionPool.close_all() clears every thread's attributes, so close() no longer found the backend PIDs of queries still running on worker threads and left them running on Redshift. RedshiftEngineAdapter now terminates running queries in both recycle() and close(), before the pool is closed. If the calling thread was interrupted mid-query or mid-transaction, its connection is closed and the sessions are terminated from a new one. Also removes the SIGTERM-to-SIGINT mapping in the CLI, which is deferred to a follow-up so this change is scoped to Ctrl-C. --- sqlmesh/cli/main.py | 8 --- sqlmesh/core/engine_adapter/redshift.py | 48 +++++++++++++----- tests/cli/test_cli.py | 20 -------- tests/core/engine_adapter/test_redshift.py | 59 +++++++++++++++++++++- 4 files changed, 91 insertions(+), 44 deletions(-) diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index b542fd7418..fc6cad8011 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -2,9 +2,7 @@ import logging import os -import signal import sys -import threading import typing as t from pathlib import Path @@ -122,12 +120,6 @@ def cli( if "--help" in sys.argv: return - if threading.current_thread() is threading.main_thread(): - # Handle SIGTERM (e.g. from CI cancellation or container shutdown) like Ctrl-C, so that the - # context is closed and running queries are cancelled before exiting. - previous_sigterm_handler = signal.signal(signal.SIGTERM, signal.default_int_handler) - ctx.call_on_close(lambda: signal.signal(signal.SIGTERM, previous_sigterm_handler)) - configure_logging( debug, log_to_stdout, diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py index 80c1a886d6..d5bf420f8e 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -144,30 +144,50 @@ 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. PG_CANCEL_BACKEND can't cancel a query inside a transaction block, so the + # 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 running_backend_pids: - if self._connection_pool.get_attribute("running_query_backend_pid") is not None: - # The calling thread was interrupted mid-query, 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) + if not running_backend_pids: + return - return super().close() + 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()) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 36e69ca4a5..f1540727b1 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -1,7 +1,6 @@ import json import os import pytest -import signal import string import time_machine from os import getcwd, path, remove @@ -2977,22 +2976,3 @@ def execute(context: ExecutionContext, **kwargs: t.Any) -> pd.DataFrame: assert result.exit_code != 0 assert "is not a known model or test file" in result.output mock.assert_not_called() - - -def test_sigterm_handled_like_keyboard_interrupt(runner, tmp_path, mocker): - create_example_project(tmp_path) - sigterm_handlers = [] - mocker.patch.object( - Context, - "plan", - side_effect=lambda *args, **kwargs: sigterm_handlers.append( - signal.getsignal(signal.SIGTERM) - ), - ) - previous_handler = signal.getsignal(signal.SIGTERM) - - result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "plan"]) - - assert result.exit_code == 0, result.output - assert sigterm_handlers == [signal.default_int_handler] - assert signal.getsignal(signal.SIGTERM) is previous_handler diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index fcd8406f89..6a18cd8731 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -962,7 +962,20 @@ def run_query(sql: str) -> None: connections[1].close.assert_called_once() -def test_close_terminates_sessions_of_queries_running_in_other_threads(mocker: MockerFixture): +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() @@ -987,7 +1000,7 @@ def run_long_query() -> None: thread = threading.Thread(target=run_long_query) thread.start() assert query_started.wait(timeout=10) - adapter.close() + teardown(adapter) thread.join(timeout=10) assert not thread.is_alive() @@ -1000,6 +1013,48 @@ def run_long_query() -> None: 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()": From 81e57b60e2ba1a35d046f5486cc2f61613332207 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Fri, 16 Oct 2026 12:23:12 -0700 Subject: [PATCH 03/10] Fix: fail transactions whose connection was closed by another thread When the thread-local connection pools close another thread's connection while it's in the middle of a transaction (e.g. when Ctrl-C recycles the pool while workers are still running), they used to forget the transaction. The worker's next statement then ran on a new connection in a new transaction and could be committed without the statements that were rolled back with the closed connection. The transaction is now marked as broken instead. The worker's next statement fails without opening a new connection, its rollback doesn't touch the database, and committing it fails. Rollback also ends the transaction now even if it fails. --- sqlmesh/utils/connection_pool.py | 52 +++++++++++-- tests/core/engine_adapter/test_redshift.py | 29 +++++++ tests/utils/test_connection_pool.py | 89 +++++++++++++++++++++- 3 files changed, 164 insertions(+), 6 deletions(-) diff --git a/sqlmesh/utils/connection_pool.py b/sqlmesh/utils/connection_pool.py index 9a70db6885..20fe7c04c0 100644 --- a/sqlmesh/utils/connection_pool.py +++ b/sqlmesh/utils/connection_pool.py @@ -3,9 +3,14 @@ import typing as t from collections import defaultdict from threading import Lock, get_ident +from threading import enumerate as enumerate_threads + +from sqlmesh.utils.errors import SQLMeshError logger = logging.getLogger(__name__) +_BROKEN_TRANSACTION_MESSAGE = "The connection was closed while a transaction was in progress, so the transaction was rolled back." + class ConnectionPool(abc.ABC): @abc.abstractmethod @@ -131,6 +136,7 @@ def __init__( self._connection_factory = connection_factory self._thread_cursors: t.Dict[t.Hashable, t.Any] = {} self._thread_transactions: t.Set[t.Hashable] = set() + self._thread_broken_transactions: t.Set[t.Hashable] = set() self._thread_attributes: t.Dict[t.Hashable, t.Dict[str, t.Any]] = defaultdict(dict) self._thread_cursors_lock = Lock() self._thread_transactions_lock = Lock() @@ -138,6 +144,7 @@ def __init__( def get_cursor(self) -> t.Any: thread_id = get_ident() + self._raise_if_transaction_broken(thread_id) with self._thread_cursors_lock: if thread_id not in self._thread_cursors: self._thread_cursors[thread_id] = self.get().cursor() @@ -167,12 +174,20 @@ def begin(self) -> None: self._thread_transactions.add(get_ident()) def commit(self) -> None: + thread_id = get_ident() + if self._is_transaction_broken(thread_id): + self._discard_transaction(thread_id) + raise SQLMeshError(_BROKEN_TRANSACTION_MESSAGE) self._do_commit() - self._discard_transaction(get_ident()) + self._discard_transaction(thread_id) def rollback(self) -> None: - self._do_rollback() - self._discard_transaction(get_ident()) + thread_id = get_ident() + try: + if not self._is_transaction_broken(thread_id): + self._do_rollback() + finally: + self._discard_transaction(thread_id) @property def is_transaction_active(self) -> bool: @@ -189,6 +204,31 @@ def close_cursor(self) -> None: def _discard_transaction(self, thread_id: t.Hashable) -> None: with self._thread_transactions_lock: self._thread_transactions.discard(thread_id) + self._thread_broken_transactions.discard(thread_id) + + def _close_transaction(self, thread_id: t.Hashable) -> None: + # The statements another thread has run in its transaction are rolled back when its + # connection is closed, so it must not continue the transaction on a new connection, where + # its remaining statements would be committed without them. The transaction is kept active + # but broken until the thread rolls it back. A finished thread can't continue it, and its + # ident may be reused by a new thread, so its transaction is discarded instead. + with self._thread_transactions_lock: + if ( + thread_id != get_ident() + and thread_id in self._thread_transactions + and any(thread.ident == thread_id for thread in enumerate_threads()) + ): + self._thread_broken_transactions.add(thread_id) + return + self._discard_transaction(thread_id) + + def _is_transaction_broken(self, thread_id: t.Hashable) -> bool: + with self._thread_transactions_lock: + return thread_id in self._thread_broken_transactions + + def _raise_if_transaction_broken(self, thread_id: t.Hashable) -> None: + if self._is_transaction_broken(thread_id): + raise SQLMeshError(_BROKEN_TRANSACTION_MESSAGE) class ThreadLocalConnectionPool(_ThreadLocalBase): @@ -203,6 +243,7 @@ def __init__( def get(self) -> t.Any: thread_id = get_ident() + self._raise_if_transaction_broken(thread_id) with self._thread_connections_lock: if thread_id not in self._thread_connections: self._thread_connections[thread_id] = self._connection_factory() @@ -226,7 +267,7 @@ def close_all(self, exclude_calling_thread: bool = False) -> None: _try_close(connection, "connection") self._thread_connections.pop(thread_id) self._thread_cursors.pop(thread_id, None) - self._discard_transaction(thread_id) + self._close_transaction(thread_id) self._thread_attributes.clear() @@ -242,6 +283,7 @@ def __init__( self._connection_lock = Lock() def get(self) -> t.Any: + self._raise_if_transaction_broken(get_ident()) with self._connection_lock: if self._connection is None: self._connection = self._connection_factory() @@ -263,7 +305,7 @@ def close_all(self, exclude_calling_thread: bool = False) -> None: if not exclude_calling_thread or thread_id != calling_thread_id: _try_close(cursor, "cursor") self._thread_cursors.pop(thread_id) - self._discard_transaction(thread_id) + self._close_transaction(thread_id) self._thread_attributes.pop(thread_id, None) if not exclude_calling_thread: diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index 6a18cd8731..af54f1d56a 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -1,6 +1,7 @@ # type: ignore import threading import typing as t +from concurrent.futures import ThreadPoolExecutor import pandas as pd # noqa: TID253 import pytest @@ -1055,6 +1056,34 @@ def run_long_query() -> None: ] +def test_recycle_fails_transaction_of_thread_between_statements(mocker: MockerFixture): + deleted = threading.Event() + recycled = threading.Event() + connect, connections = _mock_connections(mocker, lambda sql: None) + adapter = RedshiftEngineAdapter(connect, multithreaded=True) + + def run_delete_then_insert() -> None: + with pytest.raises(SQLMeshError, match="closed while a transaction was in progress"): + with adapter.transaction(): + adapter.execute("DELETE FROM t") + deleted.set() + assert recycled.wait(timeout=10) + adapter.execute("INSERT INTO t SELECT 1") + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_delete_then_insert) + assert deleted.wait(timeout=10) + adapter.recycle() + recycled.set() + future.result() + + # Closing the connection rolled back the DELETE, so the INSERT must not run and commit on a new one. + assert len(connections) == 1 + assert _executed_sql(connections[0]) == ["SELECT pg_backend_pid()", "DELETE FROM t"] + connections[0].close.assert_called_once() + connections[0].cursor.return_value.commit.assert_not_called() + + def test_close_without_running_queries(mocker: MockerFixture): def run_query(sql: str) -> None: if sql == "SELECT fail()": diff --git a/tests/utils/test_connection_pool.py b/tests/utils/test_connection_pool.py index c5926a3824..68e6f66412 100644 --- a/tests/utils/test_connection_pool.py +++ b/tests/utils/test_connection_pool.py @@ -1,6 +1,8 @@ +import typing as t from concurrent.futures import ThreadPoolExecutor -from threading import get_ident +from threading import Event, get_ident +import pytest from pytest_mock.plugin import MockerFixture from sqlmesh.utils.connection_pool import ( @@ -8,6 +10,7 @@ ThreadLocalConnectionPool, ThreadLocalSharedConnectionPool, ) +from sqlmesh.utils.errors import SQLMeshError def test_singleton_connection_pool_get(mocker: MockerFixture): @@ -323,3 +326,87 @@ def test_thread_local_shared_connection_pool_close(mocker: MockerFixture): assert cursor_mock.close.call_count == 2 assert connection_factory_mock.call_count == 2 assert connection_mock.close.call_count == 2 + + +@pytest.mark.parametrize("pool_class", [ThreadLocalConnectionPool, ThreadLocalSharedConnectionPool]) +@pytest.mark.parametrize("exclude_calling_thread", [True, False]) +def test_thread_local_pool_close_all_breaks_transactions_of_other_running_threads( + mocker: MockerFixture, pool_class: t.Type, exclude_calling_thread: bool +): + connection_factory_mock = mocker.Mock(side_effect=lambda: mocker.Mock()) + pool = pool_class(connection_factory_mock) + in_transaction = Event() + closed = Event() + + def thread() -> None: + pool.begin() + cursor = pool.get_cursor() + in_transaction.set() + assert closed.wait(timeout=10) + + # The rest of the transaction must not run on a new connection, where it would be committed + # without the statements that ran on the closed one. + assert pool.is_transaction_active + with pytest.raises(SQLMeshError, match="closed while a transaction was in progress"): + pool.get_cursor() + with pytest.raises(SQLMeshError, match="closed while a transaction was in progress"): + pool.get() + + connections_opened = connection_factory_mock.call_count + pool.rollback() + assert connection_factory_mock.call_count == connections_opened + cursor.rollback.assert_not_called() + assert not pool.is_transaction_active + + pool.get_cursor() + + pool.begin() + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(thread) + assert in_transaction.wait(timeout=10) + pool.close_all(exclude_calling_thread=exclude_calling_thread) + closed.set() + future.result() + + assert pool.is_transaction_active == exclude_calling_thread + pool.get_cursor() + + +@pytest.mark.parametrize("pool_class", [ThreadLocalConnectionPool, ThreadLocalSharedConnectionPool]) +def test_thread_local_pool_failed_rollback_ends_transaction( + mocker: MockerFixture, pool_class: t.Type +): + pool = pool_class(lambda: mocker.Mock()) + pool.begin() + pool.get_cursor().rollback.side_effect = RuntimeError("connection closed") + + with pytest.raises(RuntimeError, match="connection closed"): + pool.rollback() + assert not pool.is_transaction_active + + +@pytest.mark.parametrize("pool_class", [ThreadLocalConnectionPool, ThreadLocalSharedConnectionPool]) +def test_thread_local_pool_commit_of_broken_transaction_fails( + mocker: MockerFixture, pool_class: t.Type +): + pool = pool_class(lambda: mocker.Mock()) + in_transaction = Event() + closed = Event() + + def thread() -> None: + pool.begin() + cursor = pool.get_cursor() + in_transaction.set() + assert closed.wait(timeout=10) + + with pytest.raises(SQLMeshError, match="closed while a transaction was in progress"): + pool.commit() + cursor.commit.assert_not_called() + assert not pool.is_transaction_active + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(thread) + assert in_transaction.wait(timeout=10) + pool.close_all(exclude_calling_thread=True) + closed.set() + future.result() From 397c7083b58a5a6cfed1b8cf4ab8b119069ce285 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Fri, 16 Oct 2026 12:49:25 -0700 Subject: [PATCH 04/10] Fix: stop abandoned workers before they start a snapshot's transaction When SQLMesh is interrupted (e.g. by Ctrl-C), worker threads aren't waited for. A worker that hadn't started evaluating or creating a snapshot yet would carry on after teardown on a new connection and run the snapshot's statements anyway. concurrent_context() now sets an event when it's interrupted, and workers check it before a snapshot's non-transactional pre-statements and before its main transaction. Work after the main transaction commits isn't stopped, so the interval of committed data is still recorded. --- sqlmesh/core/snapshot/evaluator.py | 31 +++++++++ tests/core/test_snapshot_evaluator.py | 96 +++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 9d9ecf7672..192f4dfa9e 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -29,6 +29,7 @@ from collections import defaultdict from contextlib import contextmanager from functools import reduce +from threading import Event from sqlglot import exp, select from sqlglot.executor import execute @@ -146,6 +147,7 @@ def __init__( ) self.selected_gateway = selected_gateway self.ddl_concurrent_tasks = ddl_concurrent_tasks + self._interrupted = Event() def evaluate( self, @@ -662,8 +664,17 @@ def audit( @contextmanager def concurrent_context(self) -> t.Iterator[None]: + # Each context gets its own event, so that workers abandoned by an interrupted context stay + # stopped even if a new context starts before they finish. + interrupted = self._interrupted = Event() try: yield + except BaseException as ex: + if not isinstance(ex, Exception): + # Unlike failures, interrupts (e.g. KeyboardInterrupt) don't wait for worker threads + # to finish, so they're told to stop before they start running a snapshot's statements. + interrupted.set() + raise finally: self.recycle() @@ -727,6 +738,7 @@ def _evaluate_snapshot( if not snapshot.is_model: return None + interrupted = self._interrupted model = snapshot.model logger.info("Evaluating snapshot %s", snapshot.snapshot_id) @@ -772,11 +784,13 @@ def _evaluate_snapshot( ) evaluation_strategy = _evaluation_strategy(snapshot, adapter) + _raise_if_interrupted(interrupted, snapshot) evaluation_strategy.run_pre_statements( snapshot=snapshot, render_kwargs={**render_statements_kwargs, "inside_transaction": False}, ) + _raise_if_interrupted(interrupted, snapshot) with ( adapter.transaction(), adapter.session(snapshot.model.render_session_properties(**render_statements_kwargs)), @@ -888,6 +902,7 @@ def create_snapshot( if not snapshot.is_model: return + interrupted = self._interrupted logger.info("Creating a physical table for snapshot %s", snapshot.snapshot_id) adapter = self.get_adapter(snapshot.model.gateway) @@ -899,10 +914,12 @@ def create_snapshot( ) evaluation_strategy = _evaluation_strategy(snapshot, adapter) + _raise_if_interrupted(interrupted, snapshot) evaluation_strategy.run_pre_statements( snapshot=snapshot, render_kwargs={**create_render_kwargs, "inside_transaction": False} ) + _raise_if_interrupted(interrupted, snapshot) with ( adapter.transaction(), adapter.session(snapshot.model.render_session_properties(**create_render_kwargs)), @@ -3341,6 +3358,20 @@ def _intervals(snapshot: Snapshot, deployability_index: DeployabilityIndex) -> I ) +def _raise_if_interrupted(interrupted: Event, snapshot: Snapshot) -> None: + # Only checked before the outer pre-statements and before the main transaction: + # - Before these points, the worker has only read from the database (e.g. table_exists). + # - The pre-statements and the main transaction each run in a single transaction. If the worker + # is inside one when SQLMesh is interrupted, it stops anyway: its running query is terminated, + # or its next statement fails because its connection was closed. + # - After the main transaction commits, the data is written, so the worker must keep going and + # record the interval. + if interrupted.is_set(): + raise SQLMeshError( + f"Stopped processing snapshot {snapshot.snapshot_id} because SQLMesh was interrupted." + ) + + def _check_destructive_schema_change( snapshot: Snapshot, alter_operations: t.List[TableAlterOperation], diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index a9ca86496a..8436eb0728 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -1,5 +1,7 @@ from __future__ import annotations +import threading import typing as t +from concurrent.futures import Future, ThreadPoolExecutor from typing_extensions import Self from unittest.mock import call, patch, Mock @@ -5766,3 +5768,97 @@ def test_promote_resolves_this_model_with_single_mapping_entry( assert spy.call_count == 20 for call in spy.call_args_list: assert len(call.args[1]) == 1 + + +def _interrupt_while_worker_is_blocked( + evaluator: SnapshotEvaluator, blocking_mock: Mock, work: t.Callable[[], t.Any] +) -> Future: + """Runs `work` in a worker thread started inside `concurrent_context()`, interrupts the context + with a KeyboardInterrupt while the worker is blocked in `blocking_mock`, then unblocks it.""" + blocked = threading.Event() + interrupted = threading.Event() + + def block(*args: t.Any, **kwargs: t.Any) -> t.Any: + blocked.set() + assert interrupted.wait(timeout=10) + + blocking_mock.side_effect = block + + executor = ThreadPoolExecutor(max_workers=1) + try: + with pytest.raises(KeyboardInterrupt): + with evaluator.concurrent_context(): + future = executor.submit(work) + assert blocked.wait(timeout=10) + raise KeyboardInterrupt + interrupted.set() + future.exception(timeout=10) + finally: + executor.shutdown(wait=True) + return future + + +def test_evaluate_stops_before_pre_statements_when_interrupted( + adapter_mock, make_snapshot, date_kwargs: t.Dict[str, str] +): + evaluator = SnapshotEvaluator(adapter_mock) + snapshot = make_snapshot( + SqlModel(name="db.model", kind=FullKind(), query=parse_one("SELECT 1")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + future = _interrupt_while_worker_is_blocked( + evaluator, + adapter_mock.table_exists, + lambda: evaluator.evaluate(snapshot, snapshots={}, **date_kwargs), + ) + + with pytest.raises(SQLMeshError, match="interrupted"): + future.result() + adapter_mock.execute.assert_not_called() + adapter_mock.transaction.assert_not_called() + + +def test_create_snapshot_stops_before_transaction_when_interrupted(adapter_mock, make_snapshot): + evaluator = SnapshotEvaluator(adapter_mock) + snapshot = make_snapshot( + SqlModel(name="db.model", kind=FullKind(), query=parse_one("SELECT 1")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + future = _interrupt_while_worker_is_blocked( + evaluator, + adapter_mock.execute, + lambda: evaluator.create_snapshot( + snapshot, + snapshots={}, + deployability_index=DeployabilityIndex.all_deployable(), + allow_destructive_snapshots=set(), + allow_additive_snapshots=set(), + ), + ) + + with pytest.raises(SQLMeshError, match="interrupted"): + future.result() + adapter_mock.transaction.assert_not_called() + adapter_mock.create_table.assert_not_called() + + +def test_evaluate_runs_after_previous_run_was_interrupted( + adapter_mock, make_snapshot, date_kwargs: t.Dict[str, str] +): + evaluator = SnapshotEvaluator(adapter_mock) + snapshot = make_snapshot( + SqlModel(name="db.model", kind=FullKind(), query=parse_one("SELECT 1")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + with pytest.raises(KeyboardInterrupt): + with evaluator.concurrent_context(): + raise KeyboardInterrupt + + with evaluator.concurrent_context(): + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(evaluator.evaluate, snapshot, snapshots={}, **date_kwargs).result() + + adapter_mock.transaction.assert_called() From 2fba8563f29b8d2afc0a29ec7b52790659145b61 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Fri, 16 Oct 2026 16:10:35 -0700 Subject: [PATCH 05/10] Fix: stop running workers and wait for them when interrupted On interrupt, the thread pools now cancel pending tasks, call an `on_interrupt` callback to stop the running ones, and wait for them to finish before the interrupt propagates, like `ThreadPoolExecutor` does. The snapshot evaluator and scheduler pass `SnapshotEvaluator.stop_workers`, which sets the interrupted flag and recycles the adapters. Running tasks are counted rather than waited for by joining threads, since an interrupted `Thread.join` can mark a running thread as stopped (CPython < 3.13) and a thread whose start was interrupted isn't tracked by the pool. --- sqlmesh/core/scheduler.py | 1 + sqlmesh/core/snapshot/evaluator.py | 27 ++++--- sqlmesh/utils/concurrency.py | 107 ++++++++++++++++++++------ tests/core/test_snapshot_evaluator.py | 89 +++++++++++++++------ tests/utils/test_concurrency.py | 76 +++++++++++++++--- 5 files changed, 238 insertions(+), 62 deletions(-) diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py index 5eb0ff40ff..e1aee4a0a9 100644 --- a/sqlmesh/core/scheduler.py +++ b/sqlmesh/core/scheduler.py @@ -602,6 +602,7 @@ def run_node(node: SchedulingUnit) -> None: run_node, self.max_workers, raise_on_error=False, + on_interrupt=self.snapshot_evaluator.stop_workers, ) self.console.stop_evaluation_progress(success=not errors) diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 192f4dfa9e..ab6905875b 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -331,6 +331,7 @@ def promote( on_complete=on_complete, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def demote( @@ -359,6 +360,7 @@ def demote( table_mapping=table_mapping, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def create( @@ -470,6 +472,7 @@ def _create_snapshots( ), self.ddl_concurrent_tasks, raise_on_error=False, + on_interrupt=self.stop_workers, ) if errors: raise SnapshotCreationFailedError(errors, skipped) @@ -516,6 +519,7 @@ def migrate( deployability_index, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def cleanup( @@ -562,6 +566,7 @@ def cleanup( self.ddl_concurrent_tasks, reverse_order=True, raise_on_error=False, + on_interrupt=self.stop_workers, ) if errors: errored_snapshots = "\n".join(f" {e.node.name}: {e.__cause__}" for e in errors) @@ -664,20 +669,22 @@ def audit( @contextmanager def concurrent_context(self) -> t.Iterator[None]: - # Each context gets its own event, so that workers abandoned by an interrupted context stay - # stopped even if a new context starts before they finish. - interrupted = self._interrupted = Event() + # Each context gets its own event, so that workers that weren't waited for (e.g. after a + # second interrupt) stay stopped even if a new context starts before they finish. + self._interrupted = Event() try: yield - except BaseException as ex: - if not isinstance(ex, Exception): - # Unlike failures, interrupts (e.g. KeyboardInterrupt) don't wait for worker threads - # to finish, so they're told to stop before they start running a snapshot's statements. - interrupted.set() - raise finally: self.recycle() + def stop_workers(self) -> None: + """Stops the worker threads of the current concurrent context. Meant to be passed as the + `on_interrupt` of the concurrent operations run in it.""" + # Workers stop before they start a snapshot's main transaction, and recycling the adapters + # stops the queries and transactions the workers are running. + self._interrupted.set() + self.recycle() + def recycle(self) -> None: """Closes all open connections and releases all allocated resources associated with any thread except the calling one.""" @@ -1494,6 +1501,7 @@ def _create_schema( list(unique_schemas), lambda item: _create_schema(item[0], item[1], item[2]), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def get_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter: @@ -1668,6 +1676,7 @@ def _get_data_objects_in_schema( schema=s, object_names=tables_by_schema.get(s), gateway=gateway ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) for schema, objs in zip(schema_list, results): diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index 16adc8122b..f59e72308b 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -1,7 +1,7 @@ import typing as t from concurrent.futures import Executor, Future, ThreadPoolExecutor -from contextlib import contextmanager -from threading import Lock +from threading import Condition, Lock +from types import TracebackType from sqlmesh.core.snapshot import SnapshotId, SnapshotInfoLike from sqlmesh.utils.dag import DAG @@ -31,6 +31,8 @@ class ConcurrentDAGExecutor(t.Generic[H]): raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + on_interrupt: Called when the calling thread is interrupted, to stop the running tasks + before they're waited for. """ def __init__( @@ -39,11 +41,13 @@ def __init__( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool, + on_interrupt: t.Optional[t.Callable[[], None]] = None, ): self.dag = dag self.fn = fn self.tasks_num = tasks_num self.raise_on_error = raise_on_error + self.on_interrupt = on_interrupt self._init_state() @@ -59,7 +63,7 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: if self._finished_future.done(): self._init_state() - with _thread_pool(self.tasks_num) as pool: + with _StoppableThreadPool(self.tasks_num, self.on_interrupt) as pool: with self._unprocessed_nodes_lock: self._submit_next_nodes(pool) self._finished_future.result() @@ -140,6 +144,7 @@ def concurrent_apply_to_snapshots( tasks_num: int, reverse_order: bool = False, raise_on_error: bool = True, + on_interrupt: t.Optional[t.Callable[[], None]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[SnapshotId]], t.List[SnapshotId]]: """Applies a function to the given collection of snapshots concurrently while preserving the topological order between snapshots. @@ -152,6 +157,8 @@ def concurrent_apply_to_snapshots( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + on_interrupt: Called when the calling thread is interrupted, to stop the running tasks + before they're waited for. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -173,6 +180,7 @@ def concurrent_apply_to_snapshots( lambda s_id: fn(snapshots_by_id[s_id]), tasks_num, raise_on_error=raise_on_error, + on_interrupt=on_interrupt, ) @@ -181,6 +189,7 @@ def concurrent_apply_to_dag( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool = True, + on_interrupt: t.Optional[t.Callable[[], None]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: """Applies a function to the given DAG concurrently while preserving the topological order between snapshots. @@ -192,6 +201,8 @@ def concurrent_apply_to_dag( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + on_interrupt: Called when the calling thread is interrupted, to stop the running tasks + before they're waited for. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -210,6 +221,7 @@ def concurrent_apply_to_dag( fn, tasks_num, raise_on_error, + on_interrupt, ).run() @@ -250,6 +262,7 @@ def concurrent_apply_to_values( values: t.Sequence[A], fn: t.Callable[[A], R], tasks_num: int, + on_interrupt: t.Optional[t.Callable[[], None]] = None, ) -> t.List[R]: """Applies a function to the given collection of values concurrently. @@ -257,6 +270,8 @@ def concurrent_apply_to_values( values: Target values. fn: The function that will be applied concurrently to each value. tasks_num: The number of concurrent tasks. + on_interrupt: Called when the calling thread is interrupted, to stop the running tasks + before they're waited for. Returns: A list of results. @@ -272,30 +287,78 @@ def _process_value(value: A, index: int) -> None: except Exception as ex: futures[index].set_exception(ex) - with _thread_pool(tasks_num) as pool: + with _StoppableThreadPool(tasks_num, on_interrupt) 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`). +class _StoppableThreadPool(ThreadPoolExecutor): + """A `ThreadPoolExecutor` that, like its parent, waits for its tasks on exit, but handles an + interrupt of the calling thread (e.g. `KeyboardInterrupt`) differently. - 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. + When interrupted, pending tasks are cancelled and `on_interrupt`, if set, is called to stop the + running tasks (e.g. by cancelling the queries they're blocked on). The running tasks are then + waited for before the interrupt propagates. + + The running tasks are counted rather than waited for by joining the pool's threads, because a + join can't be relied on after an interrupt: + - An interrupted `Thread.join` marks the thread as stopped even if it's still running, in + CPython < 3.13 (https://github.com/python/cpython/issues/90882). + - A thread whose start was interrupted isn't tracked by the pool. `ThreadPoolExecutor` starts + a thread and only then adds it to the threads that `shutdown(wait=True)` joins: + + t.start() # `start` waits for the thread to run, so an interrupt can land here + self._threads.add(t) # skipped, though the thread may already be running a task """ - 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 + + def __init__(self, max_workers: int, on_interrupt: t.Optional[t.Callable[[], None]] = None): + super().__init__(max_workers=max_workers) + self._on_interrupt = on_interrupt + self._tasks_condition = Condition() + self._running_tasks = 0 + self._stopped = False + + def submit(self, fn: t.Callable[..., R], /, *args: t.Any, **kwargs: t.Any) -> Future: + def run() -> t.Optional[R]: + with self._tasks_condition: + if self._stopped: + return None + self._running_tasks += 1 + try: + return fn(*args, **kwargs) + finally: + with self._tasks_condition: + self._running_tasks -= 1 + self._tasks_condition.notify_all() + + return super().submit(run) + + def __exit__( + self, + exc_type: t.Optional[t.Type[BaseException]], + exc_val: t.Optional[BaseException], + exc_tb: t.Optional[TracebackType], + ) -> None: + if exc_type is None or issubclass(exc_type, Exception): + try: + self.shutdown(wait=True) + except BaseException: + self._stop() + raise + else: + self._stop() + + def _stop(self) -> None: + # Make sure no new task starts: the tasks still in the queue are cancelled, and a task that a + # worker has already taken off the queue sees `_stopped` and skips its work. + with self._tasks_condition: + self._stopped = True + self.shutdown(wait=False, cancel_futures=True) + # Ask the running tasks to stop early (e.g. by cancelling their queries). + if self._on_interrupt is not None: + self._on_interrupt() + # Wait for the running tasks to finish. + with self._tasks_condition: + self._tasks_condition.wait_for(lambda: self._running_tasks == 0) diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index 8436eb0728..d7915e3401 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -1,4 +1,5 @@ from __future__ import annotations +import signal import threading import typing as t from concurrent.futures import Future, ThreadPoolExecutor @@ -5770,31 +5771,27 @@ def test_promote_resolves_this_model_with_single_mapping_entry( assert len(call.args[1]) == 1 -def _interrupt_while_worker_is_blocked( +def _stop_workers_while_worker_is_blocked( evaluator: SnapshotEvaluator, blocking_mock: Mock, work: t.Callable[[], t.Any] ) -> Future: - """Runs `work` in a worker thread started inside `concurrent_context()`, interrupts the context - with a KeyboardInterrupt while the worker is blocked in `blocking_mock`, then unblocks it.""" + """Runs `work` in a worker thread started inside `concurrent_context()`, stops the context's + workers while the worker is blocked in `blocking_mock`, then unblocks it.""" blocked = threading.Event() - interrupted = threading.Event() + stopped = threading.Event() def block(*args: t.Any, **kwargs: t.Any) -> t.Any: blocked.set() - assert interrupted.wait(timeout=10) + assert stopped.wait(timeout=10) blocking_mock.side_effect = block - executor = ThreadPoolExecutor(max_workers=1) - try: - with pytest.raises(KeyboardInterrupt): - with evaluator.concurrent_context(): - future = executor.submit(work) - assert blocked.wait(timeout=10) - raise KeyboardInterrupt - interrupted.set() - future.exception(timeout=10) - finally: - executor.shutdown(wait=True) + with ThreadPoolExecutor(max_workers=1) as executor: + with evaluator.concurrent_context(): + future = executor.submit(work) + assert blocked.wait(timeout=10) + evaluator.stop_workers() + stopped.set() + future.exception(timeout=10) return future @@ -5807,7 +5804,7 @@ def test_evaluate_stops_before_pre_statements_when_interrupted( ) snapshot.categorize_as(SnapshotChangeCategory.BREAKING) - future = _interrupt_while_worker_is_blocked( + future = _stop_workers_while_worker_is_blocked( evaluator, adapter_mock.table_exists, lambda: evaluator.evaluate(snapshot, snapshots={}, **date_kwargs), @@ -5826,7 +5823,7 @@ def test_create_snapshot_stops_before_transaction_when_interrupted(adapter_mock, ) snapshot.categorize_as(SnapshotChangeCategory.BREAKING) - future = _interrupt_while_worker_is_blocked( + future = _stop_workers_while_worker_is_blocked( evaluator, adapter_mock.execute, lambda: evaluator.create_snapshot( @@ -5844,7 +5841,7 @@ def test_create_snapshot_stops_before_transaction_when_interrupted(adapter_mock, adapter_mock.create_table.assert_not_called() -def test_evaluate_runs_after_previous_run_was_interrupted( +def test_evaluate_runs_after_previous_run_was_stopped( adapter_mock, make_snapshot, date_kwargs: t.Dict[str, str] ): evaluator = SnapshotEvaluator(adapter_mock) @@ -5853,12 +5850,60 @@ def test_evaluate_runs_after_previous_run_was_interrupted( ) snapshot.categorize_as(SnapshotChangeCategory.BREAKING) - with pytest.raises(KeyboardInterrupt): - with evaluator.concurrent_context(): - raise KeyboardInterrupt + with evaluator.concurrent_context(): + evaluator.stop_workers() with evaluator.concurrent_context(): with ThreadPoolExecutor(max_workers=1) as executor: executor.submit(evaluator.evaluate, snapshot, snapshots={}, **date_kwargs).result() adapter_mock.transaction.assert_called() + + +def test_interrupted_create_stops_workers_and_waits_for_them( + mocker: MockerFixture, adapter_mock, make_snapshot +): + evaluator = SnapshotEvaluator(adapter_mock, ddl_concurrent_tasks=2) + snapshot = make_snapshot( + SqlModel(name="db.model", kind=FullKind(), query=parse_one("SELECT 1")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + blocked = threading.Event() + recycled = threading.Event() + worker_finished = threading.Event() + + def block_in_pre_statements(*args: t.Any, **kwargs: t.Any) -> None: + blocked.set() + assert recycled.wait(timeout=10) + + adapter_mock.execute.side_effect = block_in_pre_statements + adapter_mock.recycle.side_effect = lambda: blocked.is_set() and recycled.set() + + create_snapshot = evaluator.create_snapshot + + def create_snapshot_and_record_finish(*args: t.Any, **kwargs: t.Any) -> None: + try: + create_snapshot(*args, **kwargs) + finally: + worker_finished.set() + + mocker.patch.object(evaluator, "create_snapshot", side_effect=create_snapshot_and_record_finish) + + def interrupt_main_thread() -> None: + assert blocked.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): + evaluator.create([snapshot], {}) + # The worker was unblocked by the cancellation and waited for before the interrupt propagated. + assert worker_finished.is_set() + finally: + signal.signal(signal.SIGINT, previous_sigint_handler) + recycled.set() + + adapter_mock.transaction.assert_not_called() + adapter_mock.create_table.assert_not_called() diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 1db83c3d1d..3a0fa55e48 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,5 +1,6 @@ import signal import threading +import time import typing as t import pytest @@ -163,7 +164,7 @@ def test_concurrent_apply(tasks_num: int): ], ids=["values", "dag"], ) -def test_interrupt_does_not_wait_for_running_tasks(apply: t.Callable): +def test_interrupt_without_on_interrupt_waits_for_running_tasks(apply: t.Callable): started = [] finished = [] started_lock = threading.Lock() @@ -175,7 +176,55 @@ def fn(value: int) -> None: started.append(value) if len(started) == 2: both_started.set() - release.wait(timeout=10) + assert release.wait(timeout=10) + finished.append(value) + + def interrupt_main_thread_then_release() -> None: + both_started.wait(timeout=10) + signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) + time.sleep(0.1) + release.set() + + previous_sigint_handler = signal.signal(signal.SIGINT, signal.default_int_handler) + try: + threading.Thread(target=interrupt_main_thread_then_release).start() + with pytest.raises(KeyboardInterrupt): + apply(fn) + finally: + signal.signal(signal.SIGINT, previous_sigint_handler) + release.set() + + # The running tasks were waited for before the interrupt propagated, and the pending ones + # never started. + assert sorted(finished) == [1, 2] + assert sorted(started) == [1, 2] + + +@pytest.mark.parametrize( + "apply", + [ + lambda fn, on_interrupt: concurrent_apply_to_values( + [1, 2, 3, 4], fn, 2, on_interrupt=on_interrupt + ), + lambda fn, on_interrupt: concurrent_apply_to_dag( + DAG({1: set(), 2: set(), 3: {1}, 4: {2}}), fn, 2, on_interrupt=on_interrupt + ), + ], + ids=["values", "dag"], +) +def test_interrupt_calls_on_interrupt_then_waits_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() + assert release.wait(timeout=10) finished.append(value) def interrupt_main_thread() -> None: @@ -186,15 +235,24 @@ def interrupt_main_thread() -> None: 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 + apply(fn, release.set) 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. + # The running tasks were released by `on_interrupt` and waited for before the interrupt + # propagated, and the pending ones never started. + assert sorted(finished) == [1, 2] assert sorted(started) == [1, 2] + + +def test_error_does_not_call_on_interrupt(mocker: MockerFixture): + on_interrupt = mocker.Mock() + + def fn(value: int) -> None: + raise RuntimeError("failed") + + with pytest.raises(NodeExecutionFailedError): + concurrent_apply_to_dag(DAG({1: set()}), fn, 2, on_interrupt=on_interrupt) + + on_interrupt.assert_not_called() From b3ca68a896fefbefe46fcf892f2c92e6cf3b909b Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Fri, 16 Oct 2026 16:22:04 -0700 Subject: [PATCH 06/10] Test: skip SIGINT tests on Windows, where signal.pthread_kill isn't available --- tests/core/test_snapshot_evaluator.py | 2 ++ tests/utils/test_concurrency.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index d7915e3401..dff0bf5936 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -84,6 +84,7 @@ ) from sqlmesh.utils.metaprogramming import Executable from sqlmesh.utils.pydantic import list_of_fields_validator +from sqlmesh.utils.windows import IS_WINDOWS if t.TYPE_CHECKING: @@ -5860,6 +5861,7 @@ def test_evaluate_runs_after_previous_run_was_stopped( adapter_mock.transaction.assert_called() +@pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") def test_interrupted_create_stops_workers_and_waits_for_them( mocker: MockerFixture, adapter_mock, make_snapshot ): diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 3a0fa55e48..f193fc0884 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -14,6 +14,7 @@ concurrent_apply_to_values, ) from sqlmesh.utils.dag import DAG +from sqlmesh.utils.windows import IS_WINDOWS @pytest.mark.parametrize("tasks_num", [1, 2]) @@ -156,6 +157,7 @@ def test_concurrent_apply(tasks_num: int): assert results == [x * 2 for x in values] +@pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") @pytest.mark.parametrize( "apply", [ @@ -200,6 +202,7 @@ def interrupt_main_thread_then_release() -> None: assert sorted(started) == [1, 2] +@pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") @pytest.mark.parametrize( "apply", [ From 3c9832a8227a4211ff8d18d23d72c7c9d50df604 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Fri, 16 Oct 2026 16:52:47 -0700 Subject: [PATCH 07/10] Test: a second interrupt stops waiting for running tasks --- tests/utils/test_concurrency.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index f193fc0884..1b720f32e9 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -249,6 +249,39 @@ def interrupt_main_thread() -> None: assert sorted(started) == [1, 2] +@pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") +def test_second_interrupt_stops_waiting_for_running_tasks(): + started = threading.Semaphore(0) + ended = threading.Event() + on_interrupt_called = threading.Event() + release = threading.Event() + + def fn(value: int) -> None: + started.release() + try: + release.wait(timeout=10) + finally: + ended.set() + + def interrupt_main_thread_twice() -> None: + for _ in range(2): + assert started.acquire(timeout=10) + signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) + assert on_interrupt_called.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_twice).start() + with pytest.raises(KeyboardInterrupt): + # `on_interrupt` doesn't stop the running tasks, so only the second interrupt ends the wait. + concurrent_apply_to_values([1, 2], fn, 2, on_interrupt=on_interrupt_called.set) + assert not ended.is_set() + finally: + signal.signal(signal.SIGINT, previous_sigint_handler) + release.set() + + def test_error_does_not_call_on_interrupt(mocker: MockerFixture): on_interrupt = mocker.Mock() From 15d444ed70e8ec6376c97b43de4b8e6b8de6d0dc Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Mon, 19 Oct 2026 12:33:13 -0700 Subject: [PATCH 08/10] Feat: handle SIGTERM in the CLI like Ctrl-C CI systems cancel jobs with SIGTERM, which killed the CLI immediately and left its queries running. Raising SystemExit instead runs the same cleanup as Ctrl-C. --- sqlmesh/cli/main.py | 14 ++++++++++++++ tests/cli/test_cli.py | 16 ++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index fc6cad8011..dd4a9b66d7 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -2,7 +2,9 @@ import logging import os +import signal import sys +import threading import typing as t from pathlib import Path @@ -66,6 +68,10 @@ def _sqlmesh_version() -> str: return "0.0.0" +def _exit_on_sigterm(signum: int, frame: t.Any) -> None: + raise SystemExit(128 + signum) + + @click.group(cls=_SQLMeshGroup, no_args_is_help=True) @click.version_option(version=_sqlmesh_version(), message="%(version)s") @opt.paths @@ -120,6 +126,14 @@ def cli( if "--help" in sys.argv: return + if threading.current_thread() is threading.main_thread(): + # Handle SIGTERM (e.g. from CI cancellation) like Ctrl-C, so that running work is stopped + # and the context is closed before exiting. + previous_sigterm_handler = signal.signal(signal.SIGTERM, _exit_on_sigterm) + # The CLI can also run inside another process (e.g. tests using CliRunner), where the + # handler shouldn't outlive the command. + ctx.call_on_close(lambda: signal.signal(signal.SIGTERM, previous_sigterm_handler)) + configure_logging( debug, log_to_stdout, diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index f1540727b1..d344fab2ea 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -1,6 +1,7 @@ import json import os import pytest +import signal import string import time_machine from os import getcwd, path, remove @@ -2976,3 +2977,18 @@ def execute(context: ExecutionContext, **kwargs: t.Any) -> pd.DataFrame: assert result.exit_code != 0 assert "is not a known model or test file" in result.output mock.assert_not_called() + + +def test_sigterm_closes_context_and_exits(runner, tmp_path, mocker): + create_example_project(tmp_path) + mocker.patch.object( + Context, "plan", side_effect=lambda *args, **kwargs: signal.raise_signal(signal.SIGTERM) + ) + close = mocker.patch.object(Context, "close") + previous_handler = signal.getsignal(signal.SIGTERM) + + result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "plan"]) + + assert result.exit_code == 128 + signal.SIGTERM, result.output + close.assert_called_once() + assert signal.getsignal(signal.SIGTERM) is previous_handler From 654a3481084ead0476878cf22607bca3aeec172c Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Mon, 19 Oct 2026 17:09:06 -0700 Subject: [PATCH 09/10] Fix: only cancel queries when stopping workers, without closing their connections Stopping workers recycled the adapters, which closed the workers' connections while they were still running queries. On DuckDB 1.3, closing a connection mid-query makes the query never return, so Ctrl-C hung. Adapters now cancel their queries through cancel_queries(), which only Redshift implements. --- sqlmesh/core/engine_adapter/base.py | 4 ++++ sqlmesh/core/engine_adapter/redshift.py | 5 +++++ sqlmesh/core/snapshot/evaluator.py | 10 +++++++--- tests/core/engine_adapter/test_redshift.py | 7 +++++-- tests/core/test_snapshot_evaluator.py | 19 +++++++++++++++---- 5 files changed, 36 insertions(+), 9 deletions(-) diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 930cdf7cd4..79bb14112e 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -449,6 +449,10 @@ def recycle(self) -> None: except the calling one.""" self._connection_pool.close_all(exclude_calling_thread=True) + def cancel_queries(self) -> None: + """Stops the queries that other threads are running, if the engine supports it. Otherwise they + run until they finish.""" + def close(self) -> t.Any: """Closes all open connections and releases all allocated resources.""" self._connection_pool.close_all() diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py index d5bf420f8e..7c8533bf1d 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -157,6 +157,11 @@ def recycle(self) -> None: self._terminate_running_queries() super().recycle() + def cancel_queries(self) -> None: + # Closing the other threads' connections also fails the transactions of threads that are + # between statements, which terminating the running queries alone wouldn't stop. + self.recycle() + def close(self) -> t.Any: self._terminate_running_queries() return super().close() diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index ab6905875b..e08377f113 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -680,10 +680,14 @@ def concurrent_context(self) -> t.Iterator[None]: def stop_workers(self) -> None: """Stops the worker threads of the current concurrent context. Meant to be passed as the `on_interrupt` of the concurrent operations run in it.""" - # Workers stop before they start a snapshot's main transaction, and recycling the adapters - # stops the queries and transactions the workers are running. + # Workers stop before they start a snapshot's main transaction. Their connections aren't + # closed here, since some engines (e.g. DuckDB) hang if a connection is closed mid-query. self._interrupted.set() - self.recycle() + for gateway, adapter in self.adapters.items(): + try: + adapter.cancel_queries() + except Exception: + logger.exception("Failed to cancel the queries of gateway '%s'", gateway) def recycle(self) -> None: """Closes all open connections and releases all allocated resources associated with any thread diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index af54f1d56a..e1f5bc5ee7 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -1056,7 +1056,10 @@ def run_long_query() -> None: ] -def test_recycle_fails_transaction_of_thread_between_statements(mocker: MockerFixture): +@pytest.mark.parametrize("method", ["recycle", "cancel_queries"]) +def test_recycle_and_cancel_queries_fail_transaction_of_thread_between_statements( + mocker: MockerFixture, method: str +): deleted = threading.Event() recycled = threading.Event() connect, connections = _mock_connections(mocker, lambda sql: None) @@ -1073,7 +1076,7 @@ def run_delete_then_insert() -> None: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_delete_then_insert) assert deleted.wait(timeout=10) - adapter.recycle() + getattr(adapter, method)() recycled.set() future.result() diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index dff0bf5936..1753afd3ee 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -5872,15 +5872,15 @@ def test_interrupted_create_stops_workers_and_waits_for_them( snapshot.categorize_as(SnapshotChangeCategory.BREAKING) blocked = threading.Event() - recycled = threading.Event() + cancelled = threading.Event() worker_finished = threading.Event() def block_in_pre_statements(*args: t.Any, **kwargs: t.Any) -> None: blocked.set() - assert recycled.wait(timeout=10) + assert cancelled.wait(timeout=10) adapter_mock.execute.side_effect = block_in_pre_statements - adapter_mock.recycle.side_effect = lambda: blocked.is_set() and recycled.set() + adapter_mock.cancel_queries.side_effect = cancelled.set create_snapshot = evaluator.create_snapshot @@ -5905,7 +5905,18 @@ def interrupt_main_thread() -> None: assert worker_finished.is_set() finally: signal.signal(signal.SIGINT, previous_sigint_handler) - recycled.set() + cancelled.set() adapter_mock.transaction.assert_not_called() adapter_mock.create_table.assert_not_called() + + +def test_stop_workers_cancels_queries_without_recycling(adapter_mock): + evaluator = SnapshotEvaluator(adapter_mock) + + with evaluator.concurrent_context(): + evaluator.stop_workers() + # Recycling closes the workers' connections while they may still be using them, which some + # engines (e.g. DuckDB) don't support. + adapter_mock.cancel_queries.assert_called_once() + adapter_mock.recycle.assert_not_called() From e54ede8ea15c57b08c7ec2afc597cea2e20bcefc Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Mon, 19 Oct 2026 18:32:39 -0700 Subject: [PATCH 10/10] Test: trim redundant interrupt tests --- tests/core/engine_adapter/test_redshift.py | 32 +++++---------- tests/core/test_snapshot_evaluator.py | 21 ++-------- tests/utils/test_concurrency.py | 46 ---------------------- 3 files changed, 12 insertions(+), 87 deletions(-) diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index e1f5bc5ee7..0bec554036 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -963,19 +963,9 @@ def run_query(sql: str) -> None: 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 +@pytest.mark.parametrize("method", ["close", "cancel_queries"]) +def test_close_and_cancel_queries_terminate_sessions_of_queries_running_in_other_threads( + mocker: MockerFixture, method: str ): query_started = threading.Event() session_terminated = threading.Event() @@ -1001,7 +991,7 @@ def run_long_query() -> None: thread = threading.Thread(target=run_long_query) thread.start() assert query_started.wait(timeout=10) - teardown(adapter) + getattr(adapter, method)() thread.join(timeout=10) assert not thread.is_alive() @@ -1011,7 +1001,6 @@ def run_long_query() -> None: "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( @@ -1056,12 +1045,9 @@ def run_long_query() -> None: ] -@pytest.mark.parametrize("method", ["recycle", "cancel_queries"]) -def test_recycle_and_cancel_queries_fail_transaction_of_thread_between_statements( - mocker: MockerFixture, method: str -): +def test_cancel_queries_fails_transaction_of_thread_between_statements(mocker: MockerFixture): deleted = threading.Event() - recycled = threading.Event() + cancelled = threading.Event() connect, connections = _mock_connections(mocker, lambda sql: None) adapter = RedshiftEngineAdapter(connect, multithreaded=True) @@ -1070,14 +1056,14 @@ def run_delete_then_insert() -> None: with adapter.transaction(): adapter.execute("DELETE FROM t") deleted.set() - assert recycled.wait(timeout=10) + assert cancelled.wait(timeout=10) adapter.execute("INSERT INTO t SELECT 1") with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_delete_then_insert) assert deleted.wait(timeout=10) - getattr(adapter, method)() - recycled.set() + adapter.cancel_queries() + cancelled.set() future.result() # Closing the connection rolled back the DELETE, so the INSERT must not run and commit on a new one. diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index 1753afd3ee..dd91274d4c 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -5855,16 +5855,13 @@ def test_evaluate_runs_after_previous_run_was_stopped( evaluator.stop_workers() with evaluator.concurrent_context(): - with ThreadPoolExecutor(max_workers=1) as executor: - executor.submit(evaluator.evaluate, snapshot, snapshots={}, **date_kwargs).result() + evaluator.evaluate(snapshot, snapshots={}, **date_kwargs) adapter_mock.transaction.assert_called() @pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") -def test_interrupted_create_stops_workers_and_waits_for_them( - mocker: MockerFixture, adapter_mock, make_snapshot -): +def test_interrupted_create_cancels_queries_and_stops_workers(adapter_mock, make_snapshot): evaluator = SnapshotEvaluator(adapter_mock, ddl_concurrent_tasks=2) snapshot = make_snapshot( SqlModel(name="db.model", kind=FullKind(), query=parse_one("SELECT 1")) @@ -5873,7 +5870,6 @@ def test_interrupted_create_stops_workers_and_waits_for_them( blocked = threading.Event() cancelled = threading.Event() - worker_finished = threading.Event() def block_in_pre_statements(*args: t.Any, **kwargs: t.Any) -> None: blocked.set() @@ -5882,16 +5878,6 @@ def block_in_pre_statements(*args: t.Any, **kwargs: t.Any) -> None: adapter_mock.execute.side_effect = block_in_pre_statements adapter_mock.cancel_queries.side_effect = cancelled.set - create_snapshot = evaluator.create_snapshot - - def create_snapshot_and_record_finish(*args: t.Any, **kwargs: t.Any) -> None: - try: - create_snapshot(*args, **kwargs) - finally: - worker_finished.set() - - mocker.patch.object(evaluator, "create_snapshot", side_effect=create_snapshot_and_record_finish) - def interrupt_main_thread() -> None: assert blocked.wait(timeout=10) signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) @@ -5901,12 +5887,11 @@ def interrupt_main_thread() -> None: threading.Thread(target=interrupt_main_thread).start() with pytest.raises(KeyboardInterrupt): evaluator.create([snapshot], {}) - # The worker was unblocked by the cancellation and waited for before the interrupt propagated. - assert worker_finished.is_set() finally: signal.signal(signal.SIGINT, previous_sigint_handler) cancelled.set() + adapter_mock.cancel_queries.assert_called_once() adapter_mock.transaction.assert_not_called() adapter_mock.create_table.assert_not_called() diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 1b720f32e9..4c50910be6 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,6 +1,5 @@ import signal import threading -import time import typing as t import pytest @@ -157,51 +156,6 @@ def test_concurrent_apply(tasks_num: int): assert results == [x * 2 for x in values] -@pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") -@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_without_on_interrupt_waits_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() - assert release.wait(timeout=10) - finished.append(value) - - def interrupt_main_thread_then_release() -> None: - both_started.wait(timeout=10) - signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) - time.sleep(0.1) - release.set() - - previous_sigint_handler = signal.signal(signal.SIGINT, signal.default_int_handler) - try: - threading.Thread(target=interrupt_main_thread_then_release).start() - with pytest.raises(KeyboardInterrupt): - apply(fn) - finally: - signal.signal(signal.SIGINT, previous_sigint_handler) - release.set() - - # The running tasks were waited for before the interrupt propagated, and the pending ones - # never started. - assert sorted(finished) == [1, 2] - assert sorted(started) == [1, 2] - - @pytest.mark.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") @pytest.mark.parametrize( "apply",