From 7e4f720e62a03e913e171cdce8c42b21ad4a82f6 Mon Sep 17 00:00:00 2001 From: Gandeevan Raghuraman Date: Tue, 13 Oct 2026 17:49:29 -0700 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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()