fix(redshift): terminate running queries when SQLMesh is interrupted

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,

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:

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)
 

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,22 @@ 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, and recycling the adapters
+        # stops the queries and transactions the workers are running.
+        self._interrupted.set()
+        self.recycle()
+
     def recycle(self) -> None:
         """Closes all open connections and releases all allocated resources associated with any thread
         except the calling one."""
@@ -727,6 +745,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 +791,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 +909,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 +921,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 +1501,7 @@ def _create_schema(
                 list(unique_schemas),
                 lambda item: _create_schema(item[0], item[1], item[2]),
                 self.ddl_concurrent_tasks,
+                on_interrupt=self.stop_workers,
             )
 
     def get_adapter(self, gateway: t.Optional[str] = None) -> EngineAdapter:
@@ -1651,6 +1676,7 @@ def _get_data_objects_in_schema(
                         schema=s, object_names=tables_by_schema.get(s), gateway=gateway
                     ),
                     self.ddl_concurrent_tasks,
+                    on_interrupt=self.stop_workers,
                 )
 
                 for schema, objs in zip(schema_list, results):
@@ -3341,6 +3367,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],