feat(extract-mdl): implement new function to extract used models

ibis-server/README.md

@@ -71,7 +71,7 @@ vim .env
 ```
 Install the dependencies
 ```bash
-just install
+just install && just install-core
 ```
 Run the server
 ```bash

ibis-server/app/mdl/core.py

@@ -0,0 +1,18 @@
+from functools import cache
+
+import wren_core
+
+
+@cache
+def get_session_context(
+    manifest_str: str | None, function_path: str
+) -> wren_core.SessionContext:
+    return wren_core.SessionContext(manifest_str, function_path)
+
+
+def get_extractor(manifest_str: str) -> wren_core.Extractor:
+    return wren_core.Extractor(manifest_str)
+
+
+def to_json_base64(manifest):
+    return wren_core.to_json_base64(manifest)

ibis-server/app/mdl/rewriter.py

@@ -6,7 +6,11 @@
 from loguru import logger
 
 from app.config import get_config
-from app.mdl.context import get_session_context
+from app.mdl.core import (
+    get_extractor,
+    get_session_context,
+    to_json_base64,
+)
 from app.model import InternalServerError, UnprocessableEntityError
 from app.model.data_source import DataSource
 
@@ -59,14 +63,18 @@ def __init__(self, manifest_str: str):
 
     def rewrite(self, sql: str) -> str:
         try:
+            extractor = get_extractor(self.manifest_str)
+            tables = extractor.resolve_used_table_names(sql)
+            manifest = extractor.extract_manifest(tables)
+            manifest_str = to_json_base64(manifest)
             r = httpx.request(
                 method="GET",
                 url=f"{wren_engine_endpoint}/v2/mdl/dry-plan",
                 headers={
                     "Content-Type": "application/json",
                     "Accept": "application/json",
                 },
-                content=orjson.dumps({"manifestStr": self.manifest_str, "sql": sql}),
+                content=orjson.dumps({"manifestStr": manifest_str, "sql": sql}),
             )
             return r.raise_for_status().text.replace("\n", " ")
         except httpx.ConnectError as e:
@@ -84,7 +92,11 @@ def __init__(self, manifest_str: str, function_path: str):
 
     def rewrite(self, sql: str) -> str:
         try:
-            session_context = get_session_context(self.manifest_str, self.function_path)
+            extractor = get_extractor(self.manifest_str)
+            tables = extractor.resolve_used_table_names(sql)
+            manifest = extractor.extract_manifest(tables)
+            manifest_str = to_json_base64(manifest)
+            session_context = get_session_context(manifest_str, self.function_path)
             return session_context.transform_sql(sql)
         except Exception as e:
             raise RewriteError(str(e))

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

@@ -5,6 +5,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.rewriter import Rewriter
 from app.model import (
     DryPlanDTO,
@@ -60,10 +61,8 @@ def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Respo
 
 @router.get("/{data_source}/functions")
 def functions(data_source: DataSource) -> Response:
-    from wren_core import SessionContext
-
     file_path = get_config().get_remote_function_list_path(data_source)
-    session_context = SessionContext(None, file_path)
+    session_context = get_session_context(None, file_path)
     func_list = [f.to_dict() for f in session_context.get_available_functions()]
 
     return JSONResponse(func_list)

ibis-server/docs/development.md

@@ -27,8 +27,8 @@ This installs the pre-commit hooks.
 ## Start the server
 To get the application running:
 1. Execute `just install` to install the dependencies
-2. Create a `.env` file and fill in the required environment variables (see [Environment Variables](#Environment-Variables))
-3. If you want to use `wren_core`, you need to install the core by `just install-core`. After you modify the core, you can update it by `just update-core`.
+2. Execute `just install-core` to Install the core. If you modify the core, you can update it by `just update-core`.
+3. Create a `.env` file and fill in the required environment variables (see [Environment Variables](#Environment-Variables))
 
 To start the server:
 - Execute `just run` to start the server

ibis-server/tests/mdl/test_context.py

@@ -2,7 +2,7 @@
 
 import orjson
 
-from app.mdl.context import get_session_context
+from app.mdl.core import get_session_context
 from tests.conftest import file_path