fix(wren-ai-service): timeout in sql executor #1927
Conversation
WalkthroughRenames a local return variable in SQL utility, makes the SQL executor propagate optional Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant Web as Web Service (chart / chart_adjustment)
participant Retriever as SQL Executor (DataFetcher)
participant Engine as Engine.execute_sql
Client->>Web: Request (contains SQL)
Web->>Retriever: run(sql)
Retriever->>Engine: execute_sql(sql)
Engine-->>Retriever: (has_data, data, addition{error_message?})
alt addition.error_message present
Retriever-->>Web: {results: data, error_message}
Web->>Web: set status=failed (OTHERS), add metadata (error_type, error_message)
Web-->>Client: Failure response with error metadata
else no error_message
Retriever-->>Web: {results: data}
Web->>Web: proceed to generating -> finished / NO_CHART
Web-->>Client: Chart result (success or NO_CHART failure)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
✨ Finishing Touches
🧪 Generate 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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
yichieh-lu
added
ai-env-changed
module/ai-service
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wren-ai-service/src/pipelines/generation/utils/sql.py (1)
83-86: Align Engine.execute_sql interface with its 3-value return
Engine.execute_sql in wren-ai-service/src/core/engine.py is declared to return only(success: bool, payload: Optional[Dict[str, Any]]), but its implementations and call sites unpack three values(success, data, addition). Update the base interface’s return annotation to include the third element (or switch to a NamedTuple/TypedDict) so it matches provider implementations and allawait engine.execute_sql(…)call sites.
🧹 Nitpick comments (6)
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
119-124: Defensive access onaddition+ retain correlation for tracing.After defaulting
additionto{}, this block becomes safe. Consider also populatinginvalid_generation_result["correlation_id"]with""when absent to keep a stable schema (your code already does this later).
143-162: Consistent error typing and naming.
- Naming: in dry-run,
has_datareally means “request succeeded,” not “has data.” Considerokfor clarity.- Error codes: elsewhere (chart services) all SQL executor failures map to "OTHERS". Here you emit "TIME_OUT"/"DRY_RUN"/"PREVIEW_EMPTY_DATA"/"PREVIEW_FAILED". Consider unifying or mapping to a smaller, consistent set to simplify client handling.
Proposed minimal rename (optional):
- if has_data: + if ok: valid_generation_result = {wren-ai-service/src/pipelines/retrieval/sql_executor.py (1)
40-42: Includecorrelation_idin the return; update component outputs accordingly.You already return an extra key (
error_message) not declared in@component.output_types. Addcorrelation_idnow for traceability and declare both keys for future-proofing.Apply:
- if addition.get("error_message"): - return {"results": data, "error_message": addition.get("error_message")} - return {"results": data} + if addition.get("error_message"): + return { + "results": data, + "error_message": addition.get("error_message"), + "correlation_id": addition.get("correlation_id", ""), + } + return { + "results": data, + "correlation_id": addition.get("correlation_id", ""), + }And adjust the decorator:
- @component.output_types( - results=Optional[Dict[str, Any]], - ) + @component.output_types( + results=Optional[Dict[str, Any]], + error_message=Optional[str], + correlation_id=str, + )wren-ai-service/src/web/v1/services/chart.py (1)
120-131: Surfacecorrelation_idin state and metadata for easier debugging.If you adopt the suggested change in
DataFetcher.run, include the correlation ID here.Apply:
- if error_message: + if error_message: self._chart_results[query_id] = ChartResultResponse( status="failed", error=ChartError( code="OTHERS", message=error_message, ), - trace_id=trace_id, + trace_id=trace_id, ) results["metadata"]["error_type"] = "OTHERS" results["metadata"]["error_message"] = error_message + results["metadata"]["correlation_id"] = execute_sql_result.get("correlation_id", "") return resultswren-ai-service/src/web/v1/services/chart_adjustment.py (2)
119-127: Minor: also return and carrycorrelation_idfor parity with Chart flow.If
DataFetcher.runreturnscorrelation_id, keep it inresults["metadata"]to aid support/debugging.Apply:
execute_sql_result = ( await self._pipelines["sql_executor"].run( sql=chart_adjustment_request.sql, project_id=chart_adjustment_request.project_id, ) )["execute_sql"] sql_data = execute_sql_result["results"] error_message = execute_sql_result.get("error_message", None) + correlation_id = execute_sql_result.get("correlation_id", "")
129-141: Includecorrelation_idin failure metadata.Parity with other services and quicker incident correlation.
Apply:
self._chart_adjustment_results[ query_id ] = ChartAdjustmentResultResponse( status="failed", error=ChartAdjustmentError( code="OTHERS", message=error_message, ), ) results["metadata"]["error_type"] = "OTHERS" results["metadata"]["error_message"] = error_message + results["metadata"]["correlation_id"] = correlation_id return results
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
wren-ai-service/src/pipelines/generation/utils/sql.py(2 hunks)wren-ai-service/src/pipelines/retrieval/sql_executor.py(1 hunks)wren-ai-service/src/web/v1/services/chart.py(1 hunks)wren-ai-service/src/web/v1/services/chart_adjustment.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
wren-ai-service/src/pipelines/generation/utils/sql.py (2)
wren-ai-service/src/core/engine.py (1)
execute_sql(20-27)wren-ai-service/src/providers/engine/wren.py (3)
execute_sql(26-107)execute_sql(127-175)execute_sql(246-296)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
Summary by CodeRabbit