add required ibis api

Author: goldmedal

ibis-server/app/mdl/analyzer.py

@@ -17,3 +17,16 @@ def analyze(manifest_str: str, sql: str) -> list[dict]:
         return r.json() if r.status_code == httpx.codes.OK else r.raise_for_status()
     except httpx.ConnectError as e:
         raise ConnectionError(f"Can not connect to Wren Engine: {e}") from e
+
+
+def analyze_batch(manifest_str: str, sqls: list[str]) -> list[list[dict]]:
+    try:
+        r = httpx.request(
+            method="GET",
+            url=f"{wren_engine_endpoint}/v2/analysis/sqls",
+            headers={"Content-Type": "application/json", "Accept": "application/json"},
+            content=orjson.dumps({"manifestStr": manifest_str, "sqls": sqls}),
+        )
+        return r.json() if r.status_code == httpx.codes.OK else r.raise_for_status()
+    except httpx.ConnectError as e:
+        raise ConnectionError(f"Can not connect to Wren Engine: {e}") from e

ibis-server/app/model/__init__.py

@@ -134,6 +134,11 @@ class AnalyzeSQLDTO(BaseModel):
     sql: str
 
 
+class AnalyzeSQLBatchDTO(BaseModel):
+    manifest_str: str = manifest_str_field
+    sqls: list[str]
+
+
 class DryPlanDTO(BaseModel):
     manifest_str: str = manifest_str_field
     sql: str

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

@@ -1,8 +1,8 @@
 from fastapi import APIRouter
 
 from app.logger import log_dto
-from app.mdl.analyzer import analyze
-from app.model import AnalyzeSQLDTO
+from app.mdl.analyzer import analyze, analyze_batch
+from app.model import AnalyzeSQLBatchDTO, AnalyzeSQLDTO
 
 router = APIRouter(prefix="/analysis", tags=["analysis"])
 
@@ -11,3 +11,9 @@
 @log_dto
 def analyze_sql(dto: AnalyzeSQLDTO) -> list[dict]:
     return analyze(dto.manifest_str, dto.sql)
+
+
+@router.get("/sqls")
+@log_dto
+def analyze_sql_batch(dto: AnalyzeSQLBatchDTO) -> list[list[dict]]:
+    return analyze_batch(dto.manifest_str, dto.sqls)

ibis-server/tests/routers/v2/test_analysis.py

@@ -212,6 +212,25 @@ def test_analysis_sql_order_by():
     assert result[0]["sortings"][1]["nodeLocation"] == {"line": 1, "column": 52}
 
 
+def test_analysis_sqls():
+    result = get_sql_analysis_batch(
+        {
+            "manifestStr": manifest_str,
+            "sqls": [
+                "SELECT * FROM customer",
+                "SELECT custkey, count(*) FROM customer GROUP BY 1",
+                "WITH t1 AS (SELECT * FROM customer) SELECT * FROM t1",
+                "SELECT * FROM orders WHERE orderkey = 1 UNION SELECT * FROM orders where orderkey = 2",
+            ],
+        }
+    )
+    assert len(result) == 4
+    assert len(result[0]) == 1
+    assert len(result[1]) == 1
+    assert len(result[2]) == 2
+    assert len(result[3]) == 2
+
+
 def get_sql_analysis(input_dto):
     response = client.request(
         method="GET",
@@ -223,3 +242,16 @@ def get_sql_analysis(input_dto):
     )
     assert response.status_code == 200
     return response.json()
+
+
+def get_sql_analysis_batch(input_dto):
+    response = client.request(
+        method="GET",
+        url="/v2/analysis/sqls",
+        json={
+            "manifestStr": input_dto["manifestStr"],
+            "sqls": input_dto["sqls"],
+        },
+    )
+    assert response.status_code == 200
+    return response.json()