fix(trino): Monkeypatch TrinoQuery to cancel on KeyboardInterrupt

[OlegWock]

This patch adds monkeypatching of TrinoQuery object to cancel query if exception (Exception or KeyboardInterrupt) raised during waiting for results. This allows for cancelling queries when cell execution is cancelled, in line with what we do for BigQuery.

Here is example query that takes long time to execute. You can run it in RA using "Trino Basic Auth" integration and check query status by visiting address from the integration settings in browser.

Example query

-- Table-independent CPU load test for Trino
WITH
params AS (
  SELECT
    20000000 AS n,   -- total rows to process
    3 AS reps        -- CPU multiplier
),

a AS (
  SELECT x AS ax
  FROM UNNEST(sequence(1, 10000)) AS t(x)
),
b AS (
  SELECT y AS bx
  FROM UNNEST(sequence(1, 10000)) AS t(y)
),

-- 10k x 10k = 100M ids (trimmed later)
base AS (
  SELECT
    ((ax - 1) * 10000 + bx) AS id
  FROM a
  CROSS JOIN b
),

expanded AS (
  SELECT
    id,
    r AS rep
  FROM base
  CROSS JOIN params p
  CROSS JOIN UNNEST(sequence(1, p.reps)) AS u(r)
  WHERE id <= (SELECT n FROM params)
),

cpu_work AS (
  SELECT
    id,
    rep,
    -- CPU-heavy but streaming-friendly
    (
      sin(id * 0.000001) +
      cos(rep * 0.0001) +
      ln(1 + (id % 100000)) +
      pow(((id % 9973) + 1), 0.25)
    ) AS f,
    length(regexp_replace(cast(id AS varchar), '[0-9]', '')) AS regex_cost,
    length(replace(lpad(cast(id AS varchar), 64, '0'), '0', '')) AS string_cost
  FROM expanded
)

SELECT
  count(*) AS rows_processed,
  sum(f) AS sum_f,
  avg(f) AS avg_f,
  max(f) AS max_f,
  sum(regex_cost + string_cost) AS total_string_work
FROM cpu_work
WHERE (id % 13) <> 0

Summary by CodeRabbit

• Bug Fixes
  • Trino queries now automatically cancel when errors occur during execution (including keyboard interrupts), improving query reliability and preventing stuck operations.
• Tests
  • Added unit tests covering Trino query cancellation and behavior when cancellation fails to ensure interruptions are propagated correctly.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[linear]

BLU-5525 Trino query is not cancelled when cancelling block execution

[coderabbitai]

📝 Walkthrough

Walkthrough

A new private monkeypatch was added to intercept Trino query execution and fetching. The wrappers call the original TrinoQuery.execute/fetch, and if any exception (including KeyboardInterrupt) occurs they invoke the query's cancel() before re-raising the original exception. The patch is applied during runtime initialization alongside the BigQuery patch. Unit tests validate cancellation behavior and that cancel failures do not suppress the original exception.

Sequence Diagram(s)

sequenceDiagram
    participant Caller
    participant Runtime as RuntimePatches
    participant TrinoQuery
    participant TrinoClient as HTTPClient
    participant TrinoServer

    Caller->>TrinoQuery: call execute()/fetch()
    Note right of TrinoQuery: monkeypatched wrapper installed by RuntimePatches
    TrinoQuery->>TrinoQuery: call original execute/fetch
    TrinoQuery--xCaller: exception (e.g., KeyboardInterrupt)
    Note right of TrinoQuery: wrapper catches exception
    TrinoQuery->>TrinoQuery: invoke cancel()
    TrinoQuery->>TrinoClient: send cancel request
    TrinoClient->>TrinoServer: HTTP cancel
    TrinoServer-->>TrinoClient: cancel response
    TrinoClient-->>TrinoQuery: cancel result
    TrinoQuery-->>Caller: re-raise original exception

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately describes the main change: monkeypatching TrinoQuery to cancel on KeyboardInterrupt, matching the core objective.
Linked Issues check	✅ Passed	Changes implement the linked objective [BLU-5525]: monkeypatching TrinoQuery to cancel queries on exceptions/interrupts, with tests verifying KeyboardInterrupt handling.
Out of Scope Changes check	✅ Passed	All changes are in-scope: Trino cancellation monkey-patch, related tests, and type hint refinements align with BLU-5525 objective.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

📦 Python package built successfully!

• Version: 2.1.1.dev3+73de449
• Wheel: deepnote_toolkit-2.1.1.dev3+73de449-py3-none-any.whl
• Install:

  pip install "deepnote-toolkit @ https://deepnote-staging-runtime-artifactory.s3.amazonaws.com/deepnote-toolkit-packages/2.1.1.dev3%2B73de449/deepnote_toolkit-2.1.1.dev3%2B73de449-py3-none-any.whl"

[coderabbitai]

Actionable comments posted: 3

🤖 Fix all issues with AI agents

In `@deepnote_toolkit/runtime_patches.py`:
- Line 8: The function _monkeypatch_trino_cancel_on_error lacks an explicit
return type; update its signature to include the proper return type annotation
(e.g., -> None if it returns nothing) so it follows the coding guidelines;
locate the function definition named _monkeypatch_trino_cancel_on_error and
append the correct return type annotation and adjust any internal logic/return
statements to match that type.
- Line 1: Replace legacy typing imports and annotations: remove Any, Dict, List,
Optional, Union from the import line and update all type annotations in
deepnote_toolkit/runtime_patches.py to use modern Python 3.9+ syntax (use
dict[...] instead of Dict, list[...] instead of List, X | Y instead of Union[X,
Y], and T | None instead of Optional[T]); search for usages of Any, Dict, List,
Optional, Union in the file and convert each annotation (including function
signatures and variable annotations) to the built-in generic types and PEP 604
union syntax, leaving runtime behavior unchanged.
- Around line 27-28: The except block "except (KeyboardInterrupt, Exception):
pass" in runtime_patches.py silently swallows cancel failures; replace the pass
with a debug log call that records the exception details (e.g., logger.debug or
appropriate module logger) including the exception message/stack (use
exc_info=True or equivalent) so cancel errors are visible for troubleshooting
while still not raising them; update the except handling where "except
(KeyboardInterrupt, Exception):" appears to log and move on.

[codecov]

Codecov Report

❌ Patch coverage is 90.32258% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.36%. Comparing base (a3ac920) to head (413a16b).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
deepnote_toolkit/runtime_patches.py	90.32%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #63      +/-   ##
==========================================
+ Coverage   73.28%   73.36%   +0.08%     
==========================================
  Files          93       93              
  Lines        5199     5227      +28     
  Branches      757      758       +1     
==========================================
+ Hits         3810     3835      +25     
- Misses       1147     1149       +2     
- Partials      242      243       +1     

Flag	Coverage Δ
combined	73.36% <90.32%> (+0.08%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[deepnote-bot]

🚀 Review App Deployment Started

📝 Description	🌐 Link / Info
🌍 Review application	ra-63
🔑 Sign-in URL	Click to sign-in
📊 Application logs	View logs
🔄 Actions	Click to redeploy
🚀 ArgoCD deployment	View deployment
⏰ Last deployed	2026-01-30 09:15:21 (UTC)
📜 Deployed commit	c533743aed1aad70a4e738e25fb37ec5d3b1bdfc
🛠️ Toolkit version	73de449

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@deepnote_toolkit/runtime_patches.py`:
- Line 67: Add a docstring to the function _monkeypatch_bigquery_wait_or_cancel
describing its purpose (monkey-patch BigQuery client to modify wait_or_cancel
behavior), noting that it takes no parameters and returns None, and include
brief details about side-effects (which module/class it patches and why) so it
follows project docstring guidelines for deepnote_toolkit functions.

In `@tests/unit/test_sql_execution_internal.py`:
- Around line 33-99: Add a one-line docstring and an explicit return type hint
"-> None" to each test function:
test_trino_execute_cancels_on_keyboard_interrupt,
test_trino_fetch_cancels_on_keyboard_interrupt, and
test_trino_handles_cancel_failure_gracefully. Update each function signature to
include "-> None" and add a brief docstring describing the test's intent
immediately after the def line (e.g., "Verify TrinoQuery cancels on
KeyboardInterrupt"). Ensure no behavior changes beyond the added docstrings and
type hints.

[OlegWock]

Closing in favor of #64