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/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 8f2f098d08..7c8533bf1d 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -144,6 +144,75 @@ 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 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() + + 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/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 9d9ecf7672..e08377f113 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, @@ -329,6 +331,7 @@ def promote( on_complete=on_complete, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def demote( @@ -357,6 +360,7 @@ def demote( table_mapping=table_mapping, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def create( @@ -468,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) @@ -514,6 +519,7 @@ def migrate( deployability_index, ), self.ddl_concurrent_tasks, + on_interrupt=self.stop_workers, ) def cleanup( @@ -560,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) @@ -662,11 +669,26 @@ def audit( @contextmanager def concurrent_context(self) -> t.Iterator[None]: + # 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 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. Their connections aren't + # closed here, since some engines (e.g. DuckDB) hang if a connection is closed mid-query. + self._interrupted.set() + 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 except the calling one.""" @@ -727,6 +749,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 +795,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 +913,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 +925,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)), @@ -1477,6 +1505,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: @@ -1651,6 +1680,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): @@ -3341,6 +3371,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..f59e72308b 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -1,6 +1,7 @@ import typing as t from concurrent.futures import Executor, Future, ThreadPoolExecutor -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 @@ -30,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__( @@ -38,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() @@ -58,7 +63,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 _StoppableThreadPool(self.tasks_num, self.on_interrupt) as pool: with self._unprocessed_nodes_lock: self._submit_next_nodes(pool) self._finished_future.result() @@ -139,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. @@ -151,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. @@ -172,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, ) @@ -180,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. @@ -191,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. @@ -209,6 +221,7 @@ def concurrent_apply_to_dag( fn, tasks_num, raise_on_error, + on_interrupt, ).run() @@ -249,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. @@ -256,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. @@ -271,8 +287,78 @@ 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 _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] + + +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 `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 + """ + + 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/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/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 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..0bec554036 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,181 @@ 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() + + +@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() + + 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) + getattr(adapter, method)() + 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() + + +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_cancel_queries_fails_transaction_of_thread_between_statements(mocker: MockerFixture): + deleted = threading.Event() + cancelled = 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 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) + 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. + 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..dd91274d4c 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -1,5 +1,8 @@ from __future__ import annotations +import signal +import threading import typing as t +from concurrent.futures import Future, ThreadPoolExecutor from typing_extensions import Self from unittest.mock import call, patch, Mock @@ -81,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: @@ -5766,3 +5770,138 @@ 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 _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()`, stops the context's + workers while the worker is blocked in `blocking_mock`, then unblocks it.""" + blocked = threading.Event() + stopped = threading.Event() + + def block(*args: t.Any, **kwargs: t.Any) -> t.Any: + blocked.set() + assert stopped.wait(timeout=10) + + blocking_mock.side_effect = block + + 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 + + +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 = _stop_workers_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 = _stop_workers_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_stopped( + 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 evaluator.concurrent_context(): + evaluator.stop_workers() + + with evaluator.concurrent_context(): + 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_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")) + ) + snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + + blocked = threading.Event() + cancelled = threading.Event() + + def block_in_pre_statements(*args: t.Any, **kwargs: t.Any) -> None: + blocked.set() + assert cancelled.wait(timeout=10) + + adapter_mock.execute.side_effect = block_in_pre_statements + adapter_mock.cancel_queries.side_effect = cancelled.set + + 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], {}) + 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() + + +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() diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 5e1e4326f7..4c50910be6 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,12 +1,19 @@ +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 +from sqlmesh.utils.windows import IS_WINDOWS @pytest.mark.parametrize("tasks_num", [1, 2]) @@ -147,3 +154,95 @@ 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.skipif(IS_WINDOWS, reason="signal.pthread_kill isn't available on Windows") +@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: + 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, release.set) + finally: + signal.signal(signal.SIGINT, previous_sigint_handler) + release.set() + + # 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] + + +@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() + + 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() 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()