feat(ibis): add header to disable the fallback and refine fallback log#1174
feat(ibis): add header to disable the fallback and refine fallback log#1174douenergy merged 7 commits intoCanner:mainfrom
Conversation
WalkthroughThis update introduces a controlled fallback mechanism between v3 and v2 connector endpoints in the Ibis server. A new header, Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant V3_Endpoint
participant Util
participant V2_Endpoint
Client->>V3_Endpoint: Send request (with/without X_WREN_FALLBACK_DISABLE)
V3_Endpoint->>V3_Endpoint: Process request
alt Exception occurs
V3_Endpoint->>V3_Endpoint: Check X_WREN_FALLBACK_DISABLE header
alt Fallback disabled
V3_Endpoint-->>Client: Raise original exception (error response)
else Fallback allowed
V3_Endpoint->>Util: append_fallback_context(headers, span)
V3_Endpoint->>V2_Endpoint: Call v2 endpoint with headers, is_fallback=True
V2_Endpoint->>Util: get_fallback_message(...)
V2_Endpoint-->>V3_Endpoint: Return fallback result
V3_Endpoint-->>Client: Return fallback result
end
else No exception
V3_Endpoint-->>Client: Return result
end
Possibly related PRs
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (1)
✨ Finishing Touches
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. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
ibis-server/app/routers/v3/connector.py (1)
321-326: Same boolean-parsing issue repeated — refactor to a helper.See earlier note on lines 137-143; this copy-paste carries the same exception-clobbering risk.
🧹 Nitpick comments (3)
ibis-server/app/dependencies.py (1)
26-35: Consider simplifying the header filtering logic.The current implementation has multiple if-elif branches, which could be simplified for better readability and maintainability.
-def _filter_headers(header_string: str) -> bool: - if header_string.startswith("x-wren-"): - return True - elif header_string == "traceparent": - return True - elif header_string == "tracestate": - return True - elif header_string == "sentry-trace": - return True - return False +def _filter_headers(header_string: str) -> bool: + allowed_exact_headers = {"traceparent", "tracestate", "sentry-trace"} + return header_string.startswith("x-wren-") or header_string in allowed_exact_headersibis-server/app/routers/v3/connector.py (2)
145-146: Minor logging nit — newline isn’t needed with Loguru’s lazy formatting.
logger.warning("… {}\n", str(e))adds an extra blank line to every warning.
Drop\nunless you really want the gap:-logger.warning("Failed to execute v3 query, try to fallback to v2: {}\n", str(e)) +logger.warning("Failed to execute v3 query, try to fallback to v2: {}", str(e))
163-170: Opportunity: centralise OpenTelemetry context injection.
append_fallback_context(headers, span)is repeatedly called right before delegating to v2.
Consider moving the call into the helper that builds the v2 request (or intoappend_fallback_contextitself) so each handler doesn’t need to remember this boilerplate.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
ibis-server/app/dependencies.py(1 hunks)ibis-server/app/routers/v2/connector.py(9 hunks)ibis-server/app/routers/v3/connector.py(7 hunks)ibis-server/app/util.py(2 hunks)ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py(11 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py (5)
ibis-server/tests/conftest.py (1)
client(18-23)ibis-server/tests/routers/v3/connector/postgres/test_query.py (1)
manifest_str(96-97)ibis-server/tests/routers/v3/connector/postgres/test_model_substitute.py (1)
manifest_str(46-47)ibis-server/tests/routers/v2/connector/test_postgres.py (1)
manifest_str(124-125)ibis-server/tests/routers/v3/connector/postgres/conftest.py (1)
connection_url(48-50)
ibis-server/app/dependencies.py (2)
ibis-server/app/model/data_source.py (2)
DataSource(42-67)get_dto_type(63-67)ibis-server/app/model/__init__.py (1)
QueryDTO(17-20)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: ci
🔇 Additional comments (16)
ibis-server/app/dependencies.py (2)
7-7: Good addition of the fallback disable header constant.This header constant provides a clean way to control the fallback behavior and follows good naming conventions for HTTP headers.
15-23: Good refactoring to filter header propagation.The implementation correctly filters headers to only pass those that are relevant for tracing and Wren-specific operations, which is more secure and reduces noise in the header propagation.
ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py (3)
6-6: Good import of the fallback disable header constant.Using the constant directly from the dependencies module ensures consistency across the codebase.
48-59: Well-implemented test for disabling fallback.The test correctly verifies that setting the X_WREN_FALLBACK_DISABLE header to "true" prevents fallback to v2 and returns a 422 error.
96-108: Thorough test coverage for the fallback disable feature.The tests comprehensively cover all endpoint scenarios where fallback behavior might occur, ensuring that the fallback disable header works consistently across different operations.
Also applies to: 137-149, 162-174, 209-221, 252-264, 278-290, 304-315, 328-339, 352-363
ibis-server/app/util.py (4)
10-20: Good addition of OpenTelemetry imports.The imports correctly include all necessary components for trace context propagation, which is essential for maintaining tracing across fallback scenarios.
26-27: Clear and informative migration message.The message provides helpful context for users when fallback occurs, guiding them on how to assist with the migration process.
117-128: Well-implemented fallback context propagation.The function correctly handles trace context propagation for fallback scenarios, converting spans to non-recording spans and injecting both baggage and trace context into headers.
136-146: Good implementation of fallback message logging.The function formats the message clearly, properly serializes relevant data, and logs at an appropriate level. The replacement of newlines in SQL ensures readable log formatting.
ibis-server/app/routers/v2/connector.py (7)
25-25: Good update to imports.The import statement properly includes the new utility functions needed for fallback handling.
46-47: Proper update to function parameters.The function signature now correctly accepts the filtered headers and an optional fallback flag, allowing for controlled fallback behavior.
Also applies to: 62-62
152-156: Good implementation of fallback message logging.The code correctly logs fallback events when the is_fallback flag is set, providing valuable information for debugging and monitoring.
166-172: Consistent implementation of fallback handling.The validate endpoint follows the same pattern as the query endpoint, ensuring consistent behavior and logging across different operations.
Also applies to: 186-191
245-249: Good implementation of fallback handling in dry_plan.The function correctly logs fallback events while maintaining the original functionality.
Also applies to: 257-260
269-274: Consistent implementation in dry_plan_for_data_source.The endpoint follows the same pattern as other endpoints, ensuring consistent logging and behavior.
Also applies to: 284-288
301-302: Complete coverage of all endpoints.The model_substitute endpoint includes the same fallback handling logic, ensuring that all endpoints in the v2 connector router support the new fallback mechanism.
Also applies to: 317-321
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
ibis-server/app/dependencies.py (1)
26-37: Consider optimizing the header filtering logicThe current implementation uses multiple if-elif statements which is perfectly functional but could be made more concise.
-def _filter_headers(header_string: str) -> bool: - if header_string.startswith("x-wren-"): - return True - elif header_string.startswith("x-user-"): - return True - elif header_string == "traceparent": - return True - elif header_string == "tracestate": - return True - elif header_string == "sentry-trace": - return True - return False +def _filter_headers(header_string: str) -> bool: + return ( + header_string.startswith("x-wren-") or + header_string.startswith("x-user-") or + header_string in {"traceparent", "tracestate", "sentry-trace"} + )ibis-server/app/util.py (2)
26-27: Fix grammar in migration messageThere's a minor grammatical error in the migration message.
-MIGRATION_MESSAGE = "Wren engine is migrating to Rust version now. \ - Wren AI team are appreciate if you can provide the error messages and related logs for us." +MIGRATION_MESSAGE = "Wren engine is migrating to Rust version now. \ + Wren AI team would appreciate if you can provide the error messages and related logs for us."
148-149: Consider makingsafe_strtoboolmore explicit with negative casesThe implementation handles positive cases well but is implicit about negative cases.
def safe_strtobool(val: str) -> bool: - return val.lower() in {"1", "true", "yes", "y"} + val = val.lower() if val else "" + return val in {"1", "true", "yes", "y"}This makes the handling of None values explicit and still maintains the same behavior.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
ibis-server/app/dependencies.py(1 hunks)ibis-server/app/routers/v3/connector.py(6 hunks)ibis-server/app/util.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- ibis-server/app/routers/v3/connector.py
🧰 Additional context used
🧬 Code Graph Analysis (1)
ibis-server/app/dependencies.py (2)
ibis-server/app/model/data_source.py (2)
DataSource(42-67)get_dto_type(63-67)ibis-server/app/model/__init__.py (1)
QueryDTO(17-20)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: ci
🔇 Additional comments (5)
ibis-server/app/dependencies.py (2)
7-7: LGTM: Added header constant to control fallback behaviorThe new header constant provides a clean way to allow clients to disable the v2 fallback mechanism.
15-24: LGTM: Header filtering implementationThe implementation correctly filters headers based on prefix and specific tracing headers, allowing only necessary headers to be propagated during fallback operations.
ibis-server/app/util.py (3)
10-18: LGTM: Added necessary OpenTelemetry importsThe new imports are focused on the specific OpenTelemetry components needed for context propagation, which is good practice.
117-128: LGTM: Well-implemented OpenTelemetry context propagationThe
append_fallback_contextfunction correctly handles trace context propagation:
- Handles the case when headers is None
- Creates a non-recording span to avoid duplicate spans
- Properly injects both W3C baggage and trace context into headers
This follows OpenTelemetry best practices for context propagation across service boundaries.
136-146: LGTM: Improved fallback message loggingThe implementation provides clear, structured logging for fallback events with relevant context (datasource, model hash, SQL). The SQL formatting (replacing newlines with spaces) improves log readability.
douenergy
left a comment
douenergy
left a comment
There was a problem hiding this comment.
Thanks @goldmedal, just a small nit. The rest looks good.
ibis-server/app/util.py
Outdated
| sql = sql.replace("\n", " ") | ||
|
|
||
| message = orjson.dumps( | ||
| {"datasource": datasource, "mdl_hash": mdl_hash, "sql": sql} |
There was a problem hiding this comment.
I think we encode the mdl in base64, so maybe we should just name it mdl or mdl_base64.
|
Thanks @goldmedal |
2 participants
Description
The v2 Wren core doesn't support some data sources, such as legacy MySQL (before 8). We don't need to try to fallback on this case. If we try to fallback, the error message thrown by the v2 Wren core is confusing.
New Header
Other Changes
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Tests