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/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/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/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/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..af54f1d56a 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -1,5 +1,7 @@ # type: ignore +import threading import typing as t +from concurrent.futures import ThreadPoolExecutor import pandas as pd # noqa: TID253 import pytest @@ -667,6 +669,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 +684,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 +916,192 @@ 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_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()": + 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/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() 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] 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()