chore(ibis): add the fallback mechanism for v3 API (Canner#1100)

Author: goldmedal

ibis-server/README.md

@@ -90,3 +90,28 @@ OpenTelemetry zero-code instrumentation is highly configurable. You can set the
 
 ## Contributing
 Please see [CONTRIBUTING.md](docs/CONTRIBUTING.md) for more information.
+
+### Report the Migration Issue
+Wren engine is migrating to v3 API (powered by Rust and DataFusion). However, there are some SQL issues currently.
+If you find the migration message in your log, we hope you can provide the message and related information to the Wren AI Team.
+Just raise an issue on GitHub or contact us in the Discord channel.
+
+The message would look like the following log:
+```
+2025-03-19 22:49:08.788 | [62781772-7120-4482-b7ca-4be65c8fda96] | INFO     | __init__.dispatch:14 - POST /v3/connector/postgres/query
+2025-03-19 22:49:08.788 | [62781772-7120-4482-b7ca-4be65c8fda96] | INFO     | __init__.dispatch:15 - Request params: {}
+2025-03-19 22:49:08.789 | [62781772-7120-4482-b7ca-4be65c8fda96] | INFO     | __init__.dispatch:22 - Request body: {"connectionInfo":"REDACTED","manifestStr":"eyJjYXRhbG9nIjoid3JlbiIsInNjaGVtYSI6InB1YmxpYyIsIm1vZGVscyI6W3sibmFtZSI6Im9yZGVycyIsInRhYmxlUmVmZXJlbmNlIjp7InNjaGVtYSI6InB1YmxpYyIsIm5hbWUiOiJvcmRlcnMifSwiY29sdW1ucyI6W3sibmFtZSI6Im9yZGVya2V5IiwidHlwZSI6InZhcmNoYXIiLCJleHByZXNzaW9uIjoiY2FzdChvX29yZGVya2V5IGFzIHZhcmNoYXIpIn1dfV19","sql":"SELECT orderkey FROM orders LIMIT 1"}
+2025-03-19 22:49:08.804 | [62781772-7120-4482-b7ca-4be65c8fda96] | WARN    | connector.query:61 - Failed to execute v3 query, fallback to v2: DataFusion error: ModelAnalyzeRule
+caused by
+Schema error: No field named o_orderkey.
+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.
+```
+
+#### Steps to Report an Issue
+1. **Identify the Issue**: Look for the migration message in your log files.
+2. **Gather Information**: Collect the error message and any related logs.
+3. **Report the Issue**:
+   - **GitHub**: Open an issue on our [GitHub repository](https://github.com/Canner/wren-engine/issues) and include the collected information.
+   - **Discord**: Join our [Discord channel](https://discord.gg/5DvshJqG8Z) and share the details with us.
+
+Providing detailed information helps us to diagnose and fix the issues more efficiently. Thank you for your cooperation!

ibis-server/app/routers/v2/analysis.py

@@ -6,11 +6,11 @@
 router = APIRouter(prefix="/analysis", tags=["analysis"])
 
 
-@router.get("/sql")
+@router.get("/sql", deprecated=True)
 def analyze_sql(dto: AnalyzeSQLDTO) -> list[dict]:
     return analyze(dto.manifest_str, dto.sql)
 
 
-@router.get("/sqls")
+@router.get("/sqls", deprecated=True)
 def analyze_sql_batch(dto: AnalyzeSQLBatchDTO) -> list[list[dict]]:
     return analyze_batch(dto.manifest_str, dto.sqls)

ibis-server/app/routers/v2/connector.py

@@ -30,7 +30,9 @@ def get_java_engine_connector(request: Request) -> JavaEngineConnector:
     return request.state.java_engine_connector
 
 
-@router.post("/{data_source}/query", dependencies=[Depends(verify_query_dto)])
+@router.post(
+    "/{data_source}/query", dependencies=[Depends(verify_query_dto)], deprecated=True
+)
 async def query(
     data_source: DataSource,
     dto: QueryDTO,
@@ -63,7 +65,7 @@ async def query(
         return ORJSONResponse(to_json(connector.query(rewritten_sql, limit=limit)))
 
 
-@router.post("/{data_source}/validate/{rule_name}")
+@router.post("/{data_source}/validate/{rule_name}", deprecated=True)
 async def validate(
     data_source: DataSource,
     rule_name: str,
@@ -87,7 +89,9 @@ async def validate(
         return Response(status_code=204)
 
 
-@router.post("/{data_source}/metadata/tables", response_model=list[Table])
+@router.post(
+    "/{data_source}/metadata/tables", response_model=list[Table], deprecated=True
+)
 def get_table_list(
     data_source: DataSource,
     dto: MetadataDTO,
@@ -102,7 +106,11 @@ def get_table_list(
         ).get_table_list()
 
 
-@router.post("/{data_source}/metadata/constraints", response_model=list[Constraint])
+@router.post(
+    "/{data_source}/metadata/constraints",
+    response_model=list[Constraint],
+    deprecated=True,
+)
 def get_constraints(
     data_source: DataSource,
     dto: MetadataDTO,
@@ -117,12 +125,12 @@ def get_constraints(
         ).get_constraints()
 
 
-@router.post("/{data_source}/metadata/version")
+@router.post("/{data_source}/metadata/version", deprecated=True)
 def get_db_version(data_source: DataSource, dto: MetadataDTO) -> str:
     return MetadataFactory.get_metadata(data_source, dto.connection_info).get_version()
 
 
-@router.post("/dry-plan")
+@router.post("/dry-plan", deprecated=True)
 async def dry_plan(
     dto: DryPlanDTO,
     java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
@@ -136,7 +144,7 @@ async def dry_plan(
         ).rewrite(dto.sql)
 
 
-@router.post("/{data_source}/dry-plan")
+@router.post("/{data_source}/dry-plan", deprecated=True)
 async def dry_plan_for_data_source(
     data_source: DataSource,
     dto: DryPlanDTO,
@@ -154,7 +162,7 @@ async def dry_plan_for_data_source(
         ).rewrite(dto.sql)
 
 
-@router.post("/{data_source}/model-substitute")
+@router.post("/{data_source}/model-substitute", deprecated=True)
 async def model_substitute(
     data_source: DataSource,
     dto: TranspileDTO,

ibis-server/app/routers/v3/connector.py

@@ -8,6 +8,7 @@
 from app.config import get_config
 from app.dependencies import verify_query_dto
 from app.mdl.core import get_session_context
+from app.mdl.java_engine import JavaEngineConnector
 from app.mdl.rewriter import Rewriter
 from app.mdl.substitute import ModelSubstitute
 from app.model import (
@@ -19,11 +20,16 @@
 from app.model.connector import Connector
 from app.model.data_source import DataSource
 from app.model.validator import Validator
+from app.routers import v2
+from app.routers.v2.connector import get_java_engine_connector
 from app.util import build_context, pushdown_limit, to_json
 
 router = APIRouter(prefix="/connector")
 tracer = trace.get_tracer(__name__)
 
+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."
+
 
 @router.post("/{data_source}/query", dependencies=[Depends(verify_query_dto)])
 async def query(
@@ -32,6 +38,7 @@ async def query(
     dry_run: Annotated[bool, Query(alias="dryRun")] = False,
     limit: int | None = None,
     headers: Annotated[str | None, Header()] = None,
+    java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> Response:
     span_name = (
         f"v3_query_{data_source}_dry_run" if dry_run else f"v3_query_{data_source}"
@@ -41,44 +48,68 @@ async def query(
     ):
         try:
             sql = pushdown_limit(dto.sql, limit)
-        except Exception:
-            logger.warning("Failed to pushdown limit. Using original SQL")
-            sql = dto.sql
-
-        rewritten_sql = await Rewriter(
-            dto.manifest_str, data_source=data_source, experiment=True
-        ).rewrite(sql)
-        connector = Connector(data_source, dto.connection_info)
-        if dry_run:
-            connector.dry_run(rewritten_sql)
-            return Response(status_code=204)
-        return ORJSONResponse(to_json(connector.query(rewritten_sql, limit=limit)))
+            rewritten_sql = await Rewriter(
+                dto.manifest_str, data_source=data_source, experiment=True
+            ).rewrite(sql)
+            connector = Connector(data_source, dto.connection_info)
+            if dry_run:
+                connector.dry_run(rewritten_sql)
+                return Response(status_code=204)
+            return ORJSONResponse(to_json(connector.query(rewritten_sql, limit=limit)))
+        except Exception as e:
+            logger.warning(
+                "Failed to execute v3 query, fallback to v2: {}\n" + MIGRATION_MESSAGE,
+                str(e),
+            )
+            return await v2.connector.query(
+                data_source, dto, dry_run, limit, java_engine_connector, headers
+            )
 
 
 @router.post("/dry-plan")
 async def dry_plan(
     dto: DryPlanDTO,
     headers: Annotated[str | None, Header()] = None,
+    java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> str:
     with tracer.start_as_current_span(
         name="dry_plan", kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
-        return await Rewriter(dto.manifest_str, experiment=True).rewrite(dto.sql)
+        try:
+            return await Rewriter(dto.manifest_str, experiment=True).rewrite(dto.sql)
+        except Exception as e:
+            logger.warning(
+                "Failed to execute v3 dry-plan, fallback to v2: {}\n"
+                + MIGRATION_MESSAGE,
+                str(e),
+            )
+            return await v2.connector.dry_plan(dto, java_engine_connector, headers)
 
 
 @router.post("/{data_source}/dry-plan")
 async def dry_plan_for_data_source(
     data_source: DataSource,
     dto: DryPlanDTO,
     headers: Annotated[str | None, Header()] = None,
+    java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> str:
     span_name = f"v3_dry_plan_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
-        return await Rewriter(
-            dto.manifest_str, data_source=data_source, experiment=True
-        ).rewrite(dto.sql)
+        try:
+            return await Rewriter(
+                dto.manifest_str, data_source=data_source, experiment=True
+            ).rewrite(dto.sql)
+        except Exception as e:
+            logger.warning(
+                "Failed to execute v3 dry-plan, fallback to v2: {}\n"
+                + MIGRATION_MESSAGE,
+                str(e),
+            )
+            return await v2.connector.dry_plan_for_data_source(
+                data_source, dto, java_engine_connector, headers
+            )
 
 
 @router.post("/{data_source}/validate/{rule_name}")
@@ -87,17 +118,28 @@ async def validate(
     rule_name: str,
     dto: ValidateDTO,
     headers: Annotated[str | None, Header()] = None,
+    java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> Response:
     span_name = f"v3_validate_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
-        validator = Validator(
-            Connector(data_source, dto.connection_info),
-            Rewriter(dto.manifest_str, data_source=data_source, experiment=True),
-        )
-        await validator.validate(rule_name, dto.parameters, dto.manifest_str)
-        return Response(status_code=204)
+        try:
+            validator = Validator(
+                Connector(data_source, dto.connection_info),
+                Rewriter(dto.manifest_str, data_source=data_source, experiment=True),
+            )
+            await validator.validate(rule_name, dto.parameters, dto.manifest_str)
+            return Response(status_code=204)
+        except Exception as e:
+            logger.warning(
+                "Failed to execute v3 validate, fallback to v2: {}\n"
+                + MIGRATION_MESSAGE,
+                str(e),
+            )
+            return await v2.connector.validate(
+                data_source, rule_name, dto, java_engine_connector, headers
+            )
 
 
 @router.get("/{data_source}/functions")
@@ -120,17 +162,26 @@ async def model_substitute(
     data_source: DataSource,
     dto: TranspileDTO,
     headers: Annotated[str | None, Header()] = None,
+    java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> str:
     span_name = f"v3_model-substitute_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
-        sql = ModelSubstitute(data_source, dto.manifest_str).substitute(dto.sql)
-        Connector(data_source, dto.connection_info).dry_run(
-            await Rewriter(
-                dto.manifest_str,
-                data_source=data_source,
-                experiment=True,
-            ).rewrite(sql)
-        )
-        return sql
+        try:
+            sql = ModelSubstitute(data_source, dto.manifest_str).substitute(dto.sql)
+            Connector(data_source, dto.connection_info).dry_run(
+                await Rewriter(
+                    dto.manifest_str,
+                    data_source=data_source,
+                    experiment=True,
+                ).rewrite(sql)
+            )
+            return sql
+        except Exception as e:
+            logger.warning(
+                "Failed to execute v3 model-substitute, fallback to v2: {}", str(e)
+            )
+            return await v2.connector.model_substitute(
+                data_source, dto, java_engine_connector, headers
+            )