fix: Monkeypatch BigQuery library so KeyboardInterrupt cancels active query

Author: OlegWock

deepnote_toolkit/sql/sql_execution.py

@@ -4,6 +4,7 @@
 import re
 import uuid
 import warnings
+from typing import Any
 from urllib.parse import quote
 
 import google.oauth2.credentials
@@ -24,6 +25,7 @@
     get_project_auth_headers,
 )
 from deepnote_toolkit.ipython_utils import output_sql_metadata
+from deepnote_toolkit.logging import LoggerManager
 from deepnote_toolkit.ocelots.pandas.utils import deduplicate_columns
 from deepnote_toolkit.sql.duckdb_sql import execute_duckdb_sql
 from deepnote_toolkit.sql.jinjasql_utils import render_jinja_sql_template
@@ -33,6 +35,53 @@
 from deepnote_toolkit.sql.sql_utils import is_single_select_query
 from deepnote_toolkit.sql.url_utils import replace_user_pass_in_pg_url
 
+logger = LoggerManager().get_logger()
+
+
+# TODO(BLU-5171): Temporary hack to allow cancelling BigQuery jobs on KeyboardInterrupt (e.g. when user cancels cell execution)
+# Can be removed once
+# 1. https://github.com/googleapis/python-bigquery/pull/2331 is merged and released
+# 2. Dependicies updated for the toolkit. We don't depend on google-cloud-bigquery directly, but it's transitive
+# dependency through sqlalchemy-bigquery
+def _monkeypatch_bigquery_wait_or_cancel():
+    try:
+        from typing import Optional, Union
+
+        import google.cloud.bigquery._job_helpers as _job_helpers
+        from google.cloud.bigquery import job, table
+
+        def _wait_or_cancel(
+            job_obj: job.QueryJob,
+            api_timeout: Optional[float],
+            wait_timeout: Optional[Union[object, float]],
+            retry: Optional[Any],
+            page_size: Optional[int],
+            max_results: Optional[int],
+        ) -> table.RowIterator:
+            try:
+                return job_obj.result(
+                    page_size=page_size,
+                    max_results=max_results,
+                    retry=retry,
+                    timeout=wait_timeout,
+                )
+            except (KeyboardInterrupt, Exception):
+                try:
+                    job_obj.cancel(retry=retry, timeout=api_timeout)
+                except (KeyboardInterrupt, Exception):
+                    pass
+                raise
+
+        _job_helpers._wait_or_cancel = _wait_or_cancel
+        logger.debug("Successfully monkeypatched google.cloud.bigquery._job_helpers._wait_or_cancel")
+    except ImportError:
+        logger.debug("Could not monkeypatch BigQuery _wait_or_cancel: google.cloud.bigquery not available")
+    except Exception as e:
+        logger.warning("Failed to monkeypatch BigQuery _wait_or_cancel: %s", repr(e))
+
+
+_monkeypatch_bigquery_wait_or_cancel()
+
 
 def compile_sql_query(
     skip_jinja_template_render,

tests/unit/test_sql_execution_internal.py

@@ -9,6 +9,27 @@
 from deepnote_toolkit.sql import sql_execution as se
 
 
+def test_bigquery_wait_or_cancel_handles_keyboard_interrupt():
+    import google.cloud.bigquery._job_helpers as _job_helpers
+
+    mock_job = mock.Mock()
+    mock_job.result.side_effect = KeyboardInterrupt("User interrupted")
+    mock_job.cancel = mock.Mock()
+
+    with pytest.raises(KeyboardInterrupt):
+        # _wait_or_cancel should be monkeypatched by `_monkeypatch_bigquery_wait_or_cancel`
+        _job_helpers._wait_or_cancel(
+            job_obj=mock_job,
+            api_timeout=30.0,
+            wait_timeout=60.0,
+            retry=None,
+            page_size=None,
+            max_results=None,
+        )
+
+    mock_job.cancel.assert_called_once_with(retry=None, timeout=30.0)
+
+
 def test_build_params_for_bigquery_oauth_ok():
     with mock.patch(
         "deepnote_toolkit.sql.sql_execution.bigquery.Client"