Repository navigation
Conversation
📝 WalkthroughWalkthroughThis PR adds TTL-bounded pending cancellation tracking for async DuckDB tasks, updates server query scheduling to track scheduled query IDs, and adjusts ChangesAsync query cancellation race fix
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant _process_message
participant db_async
participant run_db_task
participant cancel_query
participant handle_query_ws
_process_message->>db_async: mark_query_scheduled(query_id)
_process_message->>run_db_task: create task with query_id
cancel_query->>db_async: signal cancel_event or record pending cancellation
run_db_task->>db_async: register_query(query_id, future, cursor, cancel_event)
db_async-->>run_db_task: set cancel_event if pending cancellation exists
run_db_task->>handle_query_ws: execute or raise CancelledError
handle_query_ws->>db_async: unmark_query_scheduled(query_id)
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5c15d8afb8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| return True | ||
| return False | ||
| if not entry: | ||
| pending_cancellations[query_id] = time.monotonic() |
There was a problem hiding this comment.
Because every miss is now saved as a pending cancellation, a cancel for an ID that is not currently active (for example, a duplicate/late cancel after the prior query already finished) will abort the next query that registers the same client-supplied queryId within the 30s TTL. The WebSocket protocol accepts caller-provided IDs for correlation and cancellation, so clients that retry or rerun with the same ID can get an immediate Query was cancelled even though the new run was never cancelled; pending cancellation should be limited to IDs known to be in the scheduling gap or otherwise not applied to arbitrary future registrations.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/sqlrooms-server/sqlrooms/server/db_async.py`:
- Around line 241-247: The `_runner` flow in
`register_query`/`start_event.set()` can deadlock if `register_query(qid,
future, cursor, cancel_event)` throws before the event is released. Update the
`_runner` logic in `db_async.py` so `start_event.set()` is guaranteed to run
even when registration fails, by wrapping the registration step in a try/finally
or equivalent cleanup path. Make sure the worker can always leave
`start_event.wait()` and proceed to its cleanup/unregister logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 27795eac-9e73-473b-aed7-3a82a7c7b092
📒 Files selected for processing (2)
python/sqlrooms-server/sqlrooms/server/db_async.pypython/sqlrooms-server/sqlrooms/server/tests/test_db_async.py
ilyabo
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
python/sqlrooms-server/sqlrooms/server/server.py (1)
414-414: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueOptional: hold a reference to the created task.
asyncio.create_task(...)is fire-and-forget; the loop keeps only a weak reference, so the task can be garbage-collected mid-execution before it completes. Consider storing it in a module/connection-scoped set and discarding it viatask.add_done_callback(...). Pre-existing behavior, but this hunk touches the call site.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/sqlrooms-server/sqlrooms/server/server.py` at line 414, The fire-and-forget task created in handle_query_ws’s call site should be retained so it cannot be garbage-collected before completion. Update the asyncio.create_task(handle_query_ws(...)) usage in server.py to store the returned task in a module- or connection-scoped collection, and register a done callback to remove it when finished. Keep the change localized around handle_query_ws and the create_task call so the task lifecycle is explicitly managed.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/sqlrooms-server/sqlrooms/server/db_async.py`:
- Around line 59-62: unmark_query_scheduled currently only removes the query
from scheduled_queries, so a pre-recorded cancel can remain in
pending_cancellations and affect a later reuse of the same query_id. Update
unmark_query_scheduled to also clear the matching entry from
pending_cancellations while holding active_queries_lock, alongside the existing
scheduled_queries cleanup, so both tracking maps stay in sync.
---
Nitpick comments:
In `@python/sqlrooms-server/sqlrooms/server/server.py`:
- Line 414: The fire-and-forget task created in handle_query_ws’s call site
should be retained so it cannot be garbage-collected before completion. Update
the asyncio.create_task(handle_query_ws(...)) usage in server.py to store the
returned task in a module- or connection-scoped collection, and register a done
callback to remove it when finished. Keep the change localized around
handle_query_ws and the create_task call so the task lifecycle is explicitly
managed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 59b70818-e957-4282-b076-1db7b866c920
📒 Files selected for processing (3)
python/sqlrooms-server/sqlrooms/server/db_async.pypython/sqlrooms-server/sqlrooms/server/server.pypython/sqlrooms-server/sqlrooms/server/tests/test_db_async.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/sqlrooms-server/sqlrooms/server/tests/test_db_async.py
| def unmark_query_scheduled(query_id: str) -> None: | ||
| """Remove a query id from scheduling-gap tracking.""" | ||
| with active_queries_lock: | ||
| scheduled_queries.pop(query_id, None) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='python/sqlrooms-server/sqlrooms/server/db_async.py'
echo '--- outline ---'
ast-grep outline "$file" --view expanded || true
echo '--- relevant symbols ---'
rg -n "def (unmark_query_scheduled|register_query|handle_query_ws|run_db_task|run_duckdb)|pending_cancellations|scheduled_queries|PENDING_CANCELLATION_TTL_SECONDS|cancel_event" "$file"
echo '--- surrounding lines 1-260 ---'
sed -n '1,260p' "$file"Repository: sqlrooms/sqlrooms
Length of output: 12443
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='python/sqlrooms-server/sqlrooms/server/db_async.py'
python3 - <<'PY'
from pathlib import Path
p = Path('python/sqlrooms-server/sqlrooms/server/db_async.py')
text = p.read_text()
for needle in [
'pending_cancellations',
'def register_query',
'def unmark_query_scheduled',
'def handle_query_ws',
'def run_db_task',
'def run_duckdb',
'PENDING_CANCELLATION_TTL_SECONDS',
]:
idx = text.find(needle)
print(f'=== {needle} @ {idx} ===')
if idx != -1:
start = max(0, text.rfind('\n', 0, idx-300))
end = min(len(text), idx+1200)
print(text[start:end])
print()
PYRepository: sqlrooms/sqlrooms
Length of output: 8100
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='python/sqlrooms-server/sqlrooms/server/db_async.py'
echo '--- run_db_task section ---'
sed -n '221,340p' "$file"
echo '--- call sites for mark/unmark/register ---'
rg -n "mark_query_scheduled\(|unmark_query_scheduled\(|register_query\(" "$file" python/sqlrooms-server -g '!**/dist/**' -g '!**/build/**'Repository: sqlrooms/sqlrooms
Length of output: 5762
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "mark_query_scheduled\(|unmark_query_scheduled\(|register_query\(" python/sqlrooms-server -g '!**/dist/**' -g '!**/build/**'
echo '--- file list for likely websocket/query handlers ---'
fd -a ".*(ws|query|duckdb|db_async).*" python/sqlrooms-serverRepository: sqlrooms/sqlrooms
Length of output: 1597
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- server.py around scheduling flow ---'
sed -n '130,180p' python/sqlrooms-server/sqlrooms/server/server.py
echo
sed -n '390,430p' python/sqlrooms-server/sqlrooms/server/server.py
echo '--- db_async tests ---'
sed -n '1,220p' python/sqlrooms-server/sqlrooms/server/tests/test_db_async.py
echo '--- query.py maybe caller path ---'
sed -n '1,220p' python/sqlrooms-server/sqlrooms/server/query.pyRepository: sqlrooms/sqlrooms
Length of output: 11075
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- server.py around scheduling flow ---'
sed -n '130,180p' python/sqlrooms-server/sqlrooms/server/server.py
echo
sed -n '390,430p' python/sqlrooms-server/sqlrooms/server/server.py
echo '--- db_async tests ---'
sed -n '1,220p' python/sqlrooms-server/sqlrooms/server/tests/test_db_async.py
echo '--- query.py ---'
sed -n '1,220p' python/sqlrooms-server/sqlrooms/server/query.pyRepository: sqlrooms/sqlrooms
Length of output: 11057
Clear pending_cancellations in unmark_query_scheduled too. A cancel recorded before registration can survive this cleanup and cancel a later query that reuses the same queryId within the TTL.
🐛 Proposed fix
def unmark_query_scheduled(query_id: str) -> None:
"""Remove a query id from scheduling-gap tracking."""
with active_queries_lock:
scheduled_queries.pop(query_id, None)
+ pending_cancellations.pop(query_id, None)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def unmark_query_scheduled(query_id: str) -> None: | |
| """Remove a query id from scheduling-gap tracking.""" | |
| with active_queries_lock: | |
| scheduled_queries.pop(query_id, None) | |
| def unmark_query_scheduled(query_id: str) -> None: | |
| """Remove a query id from scheduling-gap tracking.""" | |
| with active_queries_lock: | |
| scheduled_queries.pop(query_id, None) | |
| pending_cancellations.pop(query_id, None) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/sqlrooms-server/sqlrooms/server/db_async.py` around lines 59 - 62,
unmark_query_scheduled currently only removes the query from scheduled_queries,
so a pre-recorded cancel can remain in pending_cancellations and affect a later
reuse of the same query_id. Update unmark_query_scheduled to also clear the
matching entry from pending_cancellations while holding active_queries_lock,
alongside the existing scheduled_queries cleanup, so both tracking maps stay in
sync.
Summary
Testing
Summary by CodeRabbit