fix(wren): classify bare TimeoutError as DatabaseTimeoutError in WrenEngine

[AmirF194]

Summary

WrenEngine.query()/dry_run() catch any non-WrenError exception from the connector and wrap it as WrenError(GENERIC_USER_ERROR), the same code used for a plain bad SQL query. A connector timeout gets flattened into that same generic bucket, even though four connectors (postgres, athena, canner, trino) already special-case TimeoutError in their own except clauses specifically to keep it from being treated as a generic user error.

What failure does this repair?

postgres.py's query()/dry_run() both do except (WrenError, TimeoutError): raise, so a canceled-statement TimeoutError passes through unwrapped. One layer up, engine.py's query()/dry_run() only special-case WrenError; a bare TimeoutError isn't one, so it falls into the generic except Exception branch and comes out as GENERIC_USER_ERROR, indistinguishable from a bad query. clickhouse.py doesn't hit this because its own driver-specific TIMEOUT_EXCEEDED path raises DatabaseTimeoutError (a WrenError subclass) directly, so engine.py's except WrenError: raise preserves it correctly.

Reproduced with a fake connector that raises a bare TimeoutError (mirrors what postgres.py re-raises on a canceled statement):

wren.model.error.WrenError: [GENERIC_USER_ERROR] canceling statement due to statement timeout phase=SQL_EXECUTION

After the fix, the same call raises DatabaseTimeoutError with error_code == ErrorCode.DATABASE_TIMEOUT.

How is it tested?

• tests/unit/test_engine.py::test_query_classifies_bare_timeout_as_database_timeout and ::test_dry_run_classifies_bare_timeout_as_database_timeout: both fail against unmodified engine.py (assert DatabaseTimeoutError, get WrenError[GENERIC_USER_ERROR]) and pass on this branch.
• Full tests/unit/ suite (the test-unit CI job's own invocation, --ignore=tests/unit/test_memory.py --ignore=tests/unit/test_mcp_server.py): 1078 passed, 2 skipped (pre-existing skips, unrelated to this change).
• ruff format --check src/ and ruff check src/ (the lint CI job's own commands): clean.
• Not exercised: an actual timed-out query against a live Postgres/Athena/Canner/Trino connection. The four connectors' own except (WrenError, TimeoutError): raise clauses are unchanged by this PR; only engine.py's classification of what they re-raise is new.

Duplicate check

Swept all open PRs touching core/wren/src/wren/engine.py: #2258 (adds a YTsaurus path-rewrite method after dry_run()), #2551 and #2552 (both add a basic_safety_check(sql) call inside dry_plan()). Read each diff directly; none touch the except blocks in query()/dry_run() this PR changes.

Fixes #2153

Summary by CodeRabbit

• Bug Fixes
  • Database timeout failures are now consistently reported as database timeout errors during queries and dry runs.
  • Existing Wren errors continue to be preserved, while other unexpected failures retain generic error handling.

[coderabbitai]

Walkthrough

WrenEngine.query and dry_run now convert connector TimeoutError exceptions into DatabaseTimeoutError with ErrorCode.DATABASE_TIMEOUT. Tests verify both paths and preserve exception chaining.

Changes

Database timeout handling

Layer / File(s)	Summary
Timeout classification and validation core/wren/src/wren/engine.py, core/wren/tests/unit/test_engine.py	The query and dry-run paths classify connector timeouts as DatabaseTimeoutError. Tests verify the error type and ErrorCode.DATABASE_TIMEOUT.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: goldmedal

Poem

A rabbit found a timeout in flight,
And labeled the error clear and right.
Query and dry-run now speak,
With chained causes they can seek,
While tests guard the path each night.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: classifying bare TimeoutError as DatabaseTimeoutError in WrenEngine.
Description check	✅ Passed	The description includes all required sections with a failure reproduction, actual error output, tests, CI results, and duplicate-check details.
Linked Issues check	✅ Passed	The changes address issue #2153 by distinguishing connector timeout failures from generic SQL errors in query and dry_run.
Out of Scope Changes check	✅ Passed	The code and tests remain focused on timeout classification in WrenEngine and contain no unrelated changes.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

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 `@core/wren/tests/unit/test_engine.py`:
- Around line 198-200: Update both timeout tests around the engine.query calls
and their ErrorCode.DATABASE_TIMEOUT assertions to also assert that
exc_info.value.__cause__ is a TimeoutError, verifying the original timeout is
preserved through exception chaining.

🪄 Autofix

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: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e5e0d955-9442-4a17-9708-02728c3687b1

📥 Commits

Reviewing files that changed from the base of the PR and between 1742175 and f1e1523.

📒 Files selected for processing (2)

• core/wren/src/wren/engine.py
• core/wren/tests/unit/test_engine.py

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Assert exception chaining in both timeout tests.

The tests verify ErrorCode.DATABASE_TIMEOUT, but they do not verify the preserved TimeoutError cause. Add an assertion after each error-code assertion.

Proposed assertions

         assert exc_info.value.error_code == ErrorCode.DATABASE_TIMEOUT
+        assert isinstance(exc_info.value.__cause__, TimeoutError)

Based on the PR objective, exception chaining is part of this contract.

Also applies to: 209-211

🤖 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 `@core/wren/tests/unit/test_engine.py` around lines 198 - 200, Update both
timeout tests around the engine.query calls and their ErrorCode.DATABASE_TIMEOUT
assertions to also assert that exc_info.value.__cause__ is a TimeoutError,
verifying the original timeout is preserved through exception chaining.

[goldmedal]

Thanks @AmirF194 👍