From 13fd0eeeb82365234ab2f1ca40fb0526f35d3d9d Mon Sep 17 00:00:00 2001 From: OlegWock Date: Tue, 18 Nov 2025 17:33:03 +0100 Subject: [PATCH 1/5] fix: Monkeypatch BigQuery library so KeyboardInterrupt cancels active query --- deepnote_toolkit/sql/sql_execution.py | 49 +++++++++++++++++++++++ tests/unit/test_sql_execution_internal.py | 21 ++++++++++ 2 files changed, 70 insertions(+) diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index 02892dc..d3e4680 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/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, diff --git a/tests/unit/test_sql_execution_internal.py b/tests/unit/test_sql_execution_internal.py index 9e0187c..05eecf0 100644 --- a/tests/unit/test_sql_execution_internal.py +++ b/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" From 6c748d48e2484b11cdf011f5cff9f5f657ef4924 Mon Sep 17 00:00:00 2001 From: OlegWock Date: Tue, 18 Nov 2025 17:39:34 +0100 Subject: [PATCH 2/5] Formatting --- deepnote_toolkit/sql/sql_execution.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index d3e4680..e857415 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -41,7 +41,7 @@ # 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 +# 2. Dependencies 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: @@ -73,9 +73,13 @@ def _wait_or_cancel( raise _job_helpers._wait_or_cancel = _wait_or_cancel - logger.debug("Successfully monkeypatched google.cloud.bigquery._job_helpers._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") + 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)) From d5779b02a612bb3d05f9bf9ceda3156fe34f589d Mon Sep 17 00:00:00 2001 From: Oleh Date: Tue, 18 Nov 2025 17:42:55 +0100 Subject: [PATCH 3/5] Update deepnote_toolkit/sql/sql_execution.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- deepnote_toolkit/sql/sql_execution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index e857415..8831ae5 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -43,7 +43,7 @@ # 1. https://github.com/googleapis/python-bigquery/pull/2331 is merged and released # 2. Dependencies 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(): +def _monkeypatch_bigquery_wait_or_cancel() -> None: try: from typing import Optional, Union From 098b161978433b48c3fe65f930bef18505e30dfd Mon Sep 17 00:00:00 2001 From: OlegWock Date: Wed, 19 Nov 2025 12:44:35 +0100 Subject: [PATCH 4/5] Address PR review comments --- deepnote_toolkit/runtime_initialization.py | 7 +++ deepnote_toolkit/runtime_patches.py | 53 ++++++++++++++++++++++ deepnote_toolkit/sql/sql_execution.py | 53 ---------------------- tests/unit/conftest.py | 8 ++++ 4 files changed, 68 insertions(+), 53 deletions(-) create mode 100644 deepnote_toolkit/runtime_patches.py diff --git a/deepnote_toolkit/runtime_initialization.py b/deepnote_toolkit/runtime_initialization.py index d8a7e3c..bfafaf5 100644 --- a/deepnote_toolkit/runtime_initialization.py +++ b/deepnote_toolkit/runtime_initialization.py @@ -6,6 +6,8 @@ import psycopg2.extensions import psycopg2.extras +from deepnote_toolkit.runtime_patches import apply_runtime_patches + from .dataframe_utils import add_formatters from .execute_post_start_hooks import execute_post_start_hooks from .logging import LoggerManager @@ -24,6 +26,11 @@ def init_deepnote_runtime(): logger.debug("Initializing Deepnote runtime environment started.") + try: + apply_runtime_patches() + except Exception as e: + logger.error("Failed to apply runtime patches with a error: %s", e) + # Register sparksql magic try: IPython.get_ipython().register_magics(SparkSql) diff --git a/deepnote_toolkit/runtime_patches.py b/deepnote_toolkit/runtime_patches.py new file mode 100644 index 0000000..9dda657 --- /dev/null +++ b/deepnote_toolkit/runtime_patches.py @@ -0,0 +1,53 @@ +from typing import Any, Optional, Union + +from deepnote_toolkit.logging import LoggerManager + +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. Dependencies 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: + 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.warning( + "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)) + + +def apply_runtime_patches(): + _monkeypatch_bigquery_wait_or_cancel() diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index e857415..02892dc 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -4,7 +4,6 @@ import re import uuid import warnings -from typing import Any from urllib.parse import quote import google.oauth2.credentials @@ -25,7 +24,6 @@ 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 @@ -35,57 +33,6 @@ 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. Dependencies 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, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index a871dd8..a91c999 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -7,6 +7,14 @@ import pytest +@pytest.fixture(autouse=True, scope="session") +def apply_runtime_patches() -> None: + """Apply runtime patches once before any tests run.""" + from deepnote_toolkit.runtime_patches import apply_runtime_patches + + apply_runtime_patches() + + @pytest.fixture(autouse=True) def clean_runtime_state() -> Generator[None, None, None]: """Automatically clean in-memory env state and config cache before and after each test.""" From c896b7b050b104aced307bc75443c64d0b93788a Mon Sep 17 00:00:00 2001 From: OlegWock Date: Wed, 19 Nov 2025 12:58:43 +0100 Subject: [PATCH 5/5] Rename fixture --- tests/unit/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index a91c999..5f9179d 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -8,7 +8,7 @@ @pytest.fixture(autouse=True, scope="session") -def apply_runtime_patches() -> None: +def apply_patches() -> None: """Apply runtime patches once before any tests run.""" from deepnote_toolkit.runtime_patches import apply_runtime_patches