Address PR review comments

Author: OlegWock

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)

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()

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,

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."""