diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py index 8f2f098d08..d5bf420f8e 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -144,6 +144,70 @@ def cursor(self) -> t.Any: cursor.paramstyle = "qmark" return cursor + def recycle(self) -> None: + """Closes all open connections and releases all allocated resources associated with any thread + except the calling one, after terminating the queries that are still running on any thread. + + This includes the calling thread's query, if it has one. The calling thread can't be waiting on + a query while it runs this method, so a query recorded for it was probably abandoned by a signal + (e.g. KeyboardInterrupt). That query might still running on Redshift and holding its locks, and its + connection still has unread responses, so it can't be reused. The connection is replaced instead + of kept, which leaves the calling thread with a usable connection. + """ + self._terminate_running_queries() + super().recycle() + + def close(self) -> t.Any: + self._terminate_running_queries() + return super().close() + + def _terminate_running_queries(self) -> None: + # Redshift keeps running a query after its client disconnects, holding its locks until it + # finishes, so queries that are still running (e.g. after a KeyboardInterrupt) are stopped + # explicitly. This must happen before the connection pool is closed, which discards the + # backend PIDs. PG_CANCEL_BACKEND can't cancel a query inside a transaction block, so the + # session is terminated instead, which also rolls back its transaction. + running_backend_pids = [ + pid + for pid in self._connection_pool.get_all_attributes("running_query_backend_pid") + if pid is not None + ] + if not running_backend_pids: + return + + if ( + self._connection_pool.get_attribute("running_query_backend_pid") is not None + or self._connection_pool.is_transaction_active + ): + # The calling thread was interrupted mid-query or mid-transaction, so its connection + # can't be reused. + self._connection_pool.close() + for pid in running_backend_pids: + try: + self.execute(exp.select(exp.func("pg_terminate_backend", exp.Literal.number(pid)))) + logger.info("Terminated Redshift session %s to stop its running query", pid) + except Exception: + logger.warning("Failed to terminate Redshift session %s", pid, exc_info=True) + + def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: + self._connection_pool.set_attribute("running_query_backend_pid", self._backend_pid()) + try: + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + except Exception: + self._connection_pool.set_attribute("running_query_backend_pid", None) + raise + # Not cleared on KeyboardInterrupt or SystemExit: the query is still running on the server. + self._connection_pool.set_attribute("running_query_backend_pid", None) + + def _backend_pid(self) -> t.Optional[int]: + """The process ID of the current connection's Redshift session.""" + pid = self._connection_pool.get_attribute("backend_pid") + if pid is None: + self.cursor.execute("SELECT pg_backend_pid()") + pid = self.cursor.fetchone()[0] + self._connection_pool.set_attribute("backend_pid", pid) + return pid + def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False ) -> pd.DataFrame: 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/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..6a18cd8731 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,164 @@ def test_merge_logical( 'INSERT INTO "target" ("ID", "ts") SELECT "ID", "ts" FROM (SELECT "ID" AS "ID", "ts" AS "ts", ROW_NUMBER() OVER (PARTITION BY "ID" ORDER BY "ID") AS _row_number FROM "__temp_test_abcdefgh") AS _t WHERE _row_number = 1', 'DROP TABLE IF EXISTS "__temp_test_abcdefgh"', ] + + +def _mock_connections( + mocker: MockerFixture, run_query: t.Callable[[str], None] +) -> t.Tuple[t.Callable, t.List[t.Any]]: + """A connection factory whose connections have backend PIDs 100, 101, and so on.""" + connections = [] + + def connect(): + cursor = mocker.Mock() + cursor.fetchone.return_value = (100 + len(connections),) + cursor.execute.side_effect = lambda sql, **kwargs: ( + None if sql == "SELECT pg_backend_pid()" else run_query(sql) + ) + connection = mocker.NonCallableMock() + connection.cursor.return_value = cursor + connections.append(connection) + return connection + + return connect, connections + + +def _executed_sql(connection: t.Any) -> t.List[str]: + return [call[0][0] for call in connection.cursor.return_value.execute.call_args_list] + + +def test_close_terminates_session_of_interrupted_query(mocker: MockerFixture): + def run_query(sql: str) -> None: + if sql == "SELECT long_running()": + raise KeyboardInterrupt + + connect, connections = _mock_connections(mocker, run_query) + adapter = RedshiftEngineAdapter(connect) + + with pytest.raises(KeyboardInterrupt): + adapter.execute("SELECT long_running()") + adapter.close() + + # The interrupted connection is mid-query, so the session is terminated from a new one. + connections[0].close.assert_called_once() + assert _executed_sql(connections[1]) == [ + "SELECT pg_backend_pid()", + "SELECT PG_TERMINATE_BACKEND(100)", + ] + connections[1].close.assert_called_once() + + +def _recycle_then_close(adapter: RedshiftEngineAdapter) -> None: + # On KeyboardInterrupt, SnapshotEvaluator.concurrent_context() recycles before the context closes. + adapter.recycle() + adapter.close() + + +@pytest.mark.parametrize( + "teardown", + [lambda adapter: adapter.close(), _recycle_then_close], + ids=["close", "recycle_then_close"], +) +def test_teardown_terminates_sessions_of_queries_running_in_other_threads( + mocker: MockerFixture, teardown: t.Callable +): + query_started = threading.Event() + session_terminated = threading.Event() + + def run_query(sql: str) -> None: + if sql == "SELECT long_running()": + query_started.set() + session_terminated.wait(timeout=10) + raise Exception("terminating connection due to administrator command") + if sql == "SELECT PG_TERMINATE_BACKEND(100)": + session_terminated.set() + + connect, connections = _mock_connections(mocker, run_query) + adapter = RedshiftEngineAdapter(connect, multithreaded=True) + errors = [] + + def run_long_query() -> None: + try: + adapter.execute("SELECT long_running()") + except Exception as ex: + errors.append(ex) + + thread = threading.Thread(target=run_long_query) + thread.start() + assert query_started.wait(timeout=10) + teardown(adapter) + thread.join(timeout=10) + + assert not thread.is_alive() + assert len(errors) == 1 + assert _executed_sql(connections[1]) == [ + "SELECT pg_backend_pid()", + "SELECT PG_TERMINATE_BACKEND(100)", + ] + connections[0].close.assert_called_once() + connections[1].close.assert_called_once() + + +def test_recycle_terminates_sessions_from_new_connection_after_interrupted_transaction( + mocker: MockerFixture, +): + query_started = threading.Event() + session_terminated = threading.Event() + + def run_query(sql: str) -> None: + if sql == "SELECT long_running()": + query_started.set() + session_terminated.wait(timeout=10) + raise Exception("terminating connection due to administrator command") + if sql == "SELECT PG_TERMINATE_BACKEND(100)": + session_terminated.set() + + connect, connections = _mock_connections(mocker, run_query) + adapter = RedshiftEngineAdapter(connect, multithreaded=True) + + def run_long_query() -> None: + with pytest.raises(Exception, match="terminating connection"): + adapter.execute("SELECT long_running()") + + thread = threading.Thread(target=run_long_query) + thread.start() + assert query_started.wait(timeout=10) + + with pytest.raises(KeyboardInterrupt): + with adapter.transaction(): + adapter.execute("SELECT 1") + raise KeyboardInterrupt + adapter.recycle() + thread.join(timeout=10) + + # The interrupted transaction is left open, so the session is terminated from a new connection. + assert not thread.is_alive() + connections[0].close.assert_called_once() + connections[1].close.assert_called_once() + assert _executed_sql(connections[2]) == [ + "SELECT pg_backend_pid()", + "SELECT PG_TERMINATE_BACKEND(100)", + ] + + +def test_close_without_running_queries(mocker: MockerFixture): + def run_query(sql: str) -> None: + if sql == "SELECT fail()": + raise Exception("query failed") + + connect, connections = _mock_connections(mocker, run_query) + adapter = RedshiftEngineAdapter(connect) + + adapter.execute("SELECT 1") + with pytest.raises(Exception, match="query failed"): + adapter.execute("SELECT fail()") + adapter.close() + + # The backend PID is looked up once per connection, and nothing is terminated. + assert len(connections) == 1 + assert _executed_sql(connections[0]) == [ + "SELECT pg_backend_pid()", + "SELECT 1", + "SELECT fail()", + ] + connections[0].close.assert_called_once() 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]