Provide the batch SQL analysis API

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,31 @@ 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 result[0][0]["relation"]["tableName"] == "customer"
+    assert len(result[1]) == 1
+    assert result[1][0]["relation"]["tableName"] == "customer"
+    assert len(result[2]) == 2
+    assert result[2][0]["relation"]["tableName"] == "customer"
+    assert result[2][1]["relation"]["tableName"] == "t1"
+    assert len(result[3]) == 2
+    assert result[3][0]["relation"]["tableName"] == "orders"
+    assert result[3][1]["relation"]["tableName"] == "orders"
+
+
 def get_sql_analysis(input_dto):
     response = client.request(
         method="GET",
@@ -223,3 +248,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()

wren-main/src/main/java/io/wren/main/web/AnalysisResource.java

@@ -43,6 +43,7 @@
 import static io.wren.main.web.WrenExceptionMapper.bindAsyncResponse;
 import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
 
+@Deprecated
 @Path("/v1/analysis")
 public class AnalysisResource
 {

wren-main/src/main/java/io/wren/main/web/AnalysisResourceV2.java

@@ -19,6 +19,7 @@
 import io.wren.base.SessionContext;
 import io.wren.base.WrenMDL;
 import io.wren.base.sqlrewrite.analyzer.decisionpoint.DecisionPointAnalyzer;
+import io.wren.main.web.dto.SqlAnalysisInputBatchDto;
 import io.wren.main.web.dto.SqlAnalysisInputDtoV2;
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.GET;
@@ -72,4 +73,35 @@ public void getSqlAnalysis(
                 })
                 .whenComplete(bindAsyncResponse(asyncResponse));
     }
+
+    @GET
+    @Path("/sqls")
+    @Consumes(APPLICATION_JSON)
+    @Produces(APPLICATION_JSON)
+    public void getSqlAnalysisBatch(
+            SqlAnalysisInputBatchDto inputBatchDto,
+            @Suspended AsyncResponse asyncResponse)
+    {
+        CompletableFuture
+                .supplyAsync(() ->
+                        Optional.ofNullable(inputBatchDto.getManifestStr())
+                                .orElseThrow(() -> new IllegalArgumentException("Manifest is required")))
+                .thenApply(manifestStr -> {
+                    try {
+                        return WrenMDL.fromJson(new String(Base64.getDecoder().decode(manifestStr), UTF_8));
+                    }
+                    catch (IOException e) {
+                        throw new RuntimeException(e);
+                    }
+                })
+                .thenApply(mdl ->
+                  inputBatchDto.getSqls().stream().map(sql -> {
+                      Statement statement = parseSql(sql);
+                      return DecisionPointAnalyzer.analyze(
+                            statement,
+                            SessionContext.builder().setCatalog(mdl.getCatalog()).setSchema(mdl.getSchema()).build(),
+                            mdl).stream().map(AnalysisResource::toQueryAnalysisDto).toList();
+                  }).toList())
+                .whenComplete(bindAsyncResponse(asyncResponse));
+    }
 }

wren-main/src/main/java/io/wren/main/web/dto/SqlAnalysisInputBatchDto.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.wren.main.web.dto;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.List;
+
+public class SqlAnalysisInputBatchDto
+{
+    private final String manifestStr;
+    private final List<String> sqls;
+
+    @JsonCreator
+    public SqlAnalysisInputBatchDto(String manifestStr, List<String> sqls)
+    {
+        this.manifestStr = manifestStr;
+        this.sqls = sqls == null ? List.of() : sqls;
+    }
+
+    @JsonProperty
+    public String getManifestStr()
+    {
+        return manifestStr;
+    }
+
+    @JsonProperty
+    public List<String> getSqls()
+    {
+        return sqls;
+    }
+}

wren-tests/src/test/java/io/wren/testing/RequireWrenServer.java

@@ -30,14 +30,13 @@
 import io.wren.base.dto.Manifest;
 import io.wren.main.connector.duckdb.DuckDBMetadata;
 import io.wren.main.validation.ValidationResult;
-import io.wren.main.web.dto.CheckOutputDto;
-import io.wren.main.web.dto.DeployInputDto;
 import io.wren.main.web.dto.DryPlanDto;
 import io.wren.main.web.dto.DryPlanDtoV2;
 import io.wren.main.web.dto.ErrorMessageDto;
 import io.wren.main.web.dto.PreviewDto;
 import io.wren.main.web.dto.QueryAnalysisDto;
 import io.wren.main.web.dto.QueryResultDto;
+import io.wren.main.web.dto.SqlAnalysisInputBatchDto;
 import io.wren.main.web.dto.SqlAnalysisInputDto;
 import io.wren.main.web.dto.ValidateDto;
 import jakarta.ws.rs.WebApplicationException;
@@ -48,10 +47,6 @@
 
 import java.io.IOException;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
 import static io.airlift.http.client.JsonBodyGenerator.jsonBodyGenerator;
@@ -76,10 +71,8 @@ public abstract class RequireWrenServer
     protected HttpClient client;
 
     private static final JsonCodec<ErrorMessageDto> ERROR_CODEC = jsonCodec(ErrorMessageDto.class);
-    private static final JsonCodec<CheckOutputDto> CHECK_OUTPUT_DTO_CODEC = jsonCodec(CheckOutputDto.class);
     public static final JsonCodec<Manifest> MANIFEST_JSON_CODEC = jsonCodec(Manifest.class);
     private static final JsonCodec<PreviewDto> PREVIEW_DTO_CODEC = jsonCodec(PreviewDto.class);
-    private static final JsonCodec<DeployInputDto> DEPLOY_INPUT_DTO_JSON_CODEC = jsonCodec(DeployInputDto.class);
     private static final JsonCodec<SqlAnalysisInputDto> SQL_ANALYSIS_INPUT_DTO_CODEC = jsonCodec(SqlAnalysisInputDto.class);
     private static final JsonCodec<ConfigManager.ConfigEntry> CONFIG_ENTRY_JSON_CODEC = jsonCodec(ConfigManager.ConfigEntry.class);
     private static final JsonCodec<List<ConfigManager.ConfigEntry>> CONFIG_ENTRY_LIST_CODEC = listJsonCodec(ConfigManager.ConfigEntry.class);
@@ -90,6 +83,8 @@ public abstract class RequireWrenServer
     private static final JsonCodec<List<ValidationResult>> VALIDATION_RESULT_LIST_CODEC = listJsonCodec(ValidationResult.class);
     private static final JsonCodec<ValidateDto> VALIDATE_DTO_CODEC = jsonCodec(ValidateDto.class);
     private static final JsonCodec<List<QueryAnalysisDto>> QUERY_ANALYSIS_DTO_LIST_CODEC = listJsonCodec(QueryAnalysisDto.class);
+    private static final JsonCodec<SqlAnalysisInputBatchDto> SQL_ANALYSIS_INPUT_BATCH_DTO_CODEC = jsonCodec(SqlAnalysisInputBatchDto.class);
+    private static final JsonCodec<List<List<QueryAnalysisDto>>> QUERY_ANALYSIS_DTO_LIST_LIST_CODEC = listJsonCodec(listJsonCodec(QueryAnalysisDto.class));
 
     public RequireWrenServer() {}
 
@@ -215,78 +210,34 @@ protected String dryPlanV2(DryPlanDtoV2 dryPlanDto)
         return response.getBody();
     }
 
-    protected void deployMDL(DeployInputDto dto)
-    {
-        Request request = preparePost()
-                .setUri(server().getHttpServerBasedUrl().resolve("/v1/mdl/deploy"))
-                .setHeader(CONTENT_TYPE, "application/json")
-                .setBodyGenerator(jsonBodyGenerator(DEPLOY_INPUT_DTO_JSON_CODEC, dto))
-                .build();
-
-        StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
-        if (response.getStatusCode() != 202) {
-            getWebApplicationException(response);
-        }
-    }
-
-    protected Manifest getCurrentManifest()
-    {
-        Request request = prepareGet()
-                .setUri(server().getHttpServerBasedUrl().resolve("/v1/mdl"))
-                .build();
-
-        StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
-        if (response.getStatusCode() != 200) {
-            getWebApplicationException(response);
-        }
-        return MANIFEST_JSON_CODEC.fromJson(response.getBody());
-    }
-
-    protected CheckOutputDto getDeployStatus()
+    protected List<QueryAnalysisDto> getSqlAnalysis(SqlAnalysisInputDto inputDto)
     {
         Request request = prepareGet()
-                .setUri(server().getHttpServerBasedUrl().resolve("/v1/mdl/status"))
+                .setUri(server().getHttpServerBasedUrl().resolve("/v1/analysis/sql"))
+                .setHeader(CONTENT_TYPE, "application/json")
+                .setBodyGenerator(jsonBodyGenerator(SQL_ANALYSIS_INPUT_DTO_CODEC, inputDto))
                 .build();
 
         StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
         if (response.getStatusCode() != 200) {
             getWebApplicationException(response);
         }
-        return CHECK_OUTPUT_DTO_CODEC.fromJson(response.getBody());
-    }
-
-    protected void waitUntilReady()
-            throws ExecutionException, InterruptedException, TimeoutException
-    {
-        CompletableFuture.runAsync(() -> {
-            while (true) {
-                CheckOutputDto checkOutputDto = getDeployStatus();
-                if (checkOutputDto.getStatus() == CheckOutputDto.Status.READY) {
-                    break;
-                }
-                try {
-                    Thread.sleep(1000);
-                }
-                catch (InterruptedException e) {
-                    throw new AssertionError("Status doesn't change to READY", e);
-                }
-            }
-        }).get(60, TimeUnit.SECONDS);
+        return QUERY_ANALYSIS_DTO_LIST_CODEC.fromJson(response.getBody());
     }
 
-    protected List<QueryAnalysisDto> getSqlAnalysis(SqlAnalysisInputDto inputDto)
+    protected List<List<QueryAnalysisDto>> getSqlAnalysisBatch(SqlAnalysisInputBatchDto inputBatchDto)
     {
         Request request = prepareGet()
-                .setUri(server().getHttpServerBasedUrl().resolve("/v1/analysis/sql"))
+                .setUri(server().getHttpServerBasedUrl().resolve("/v2/analysis/sqls"))
                 .setHeader(CONTENT_TYPE, "application/json")
-                .setBodyGenerator(jsonBodyGenerator(SQL_ANALYSIS_INPUT_DTO_CODEC, inputDto))
+                .setBodyGenerator(jsonBodyGenerator(SQL_ANALYSIS_INPUT_BATCH_DTO_CODEC, inputBatchDto))
                 .build();
 
         StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
         if (response.getStatusCode() != 200) {
             getWebApplicationException(response);
         }
-        return QUERY_ANALYSIS_DTO_LIST_CODEC.fromJson(response.getBody());
+        return QUERY_ANALYSIS_DTO_LIST_LIST_CODEC.fromJson(response.getBody());
     }
 
     protected List<ConfigManager.ConfigEntry> getConfigs()

wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java

@@ -22,19 +22,22 @@
 import io.wren.base.sqlrewrite.analyzer.decisionpoint.FilterAnalysis;
 import io.wren.base.sqlrewrite.analyzer.decisionpoint.RelationAnalysis;
 import io.wren.main.web.dto.QueryAnalysisDto;
+import io.wren.main.web.dto.SqlAnalysisInputBatchDto;
 import io.wren.main.web.dto.SqlAnalysisInputDto;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Base64;
 import java.util.List;
 import java.util.Set;
 
 import static io.wren.base.dto.Model.onTableReference;
 import static io.wren.base.dto.TableReference.tableReference;
 import static io.wren.main.web.dto.NodeLocationDto.nodeLocationDto;
 import static io.wren.testing.AbstractTestFramework.DEFAULT_SESSION_CONTEXT;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TestAnalysisResource
@@ -184,4 +187,32 @@ public void testBasic()
         assertThat(result.get(0).getSortings().get(1).getOrdering()).isEqualTo(SortItem.Ordering.DESCENDING.name());
         assertThat(result.get(0).getSortings().get(1).getNodeLocation()).isEqualTo(nodeLocationDto(1, 52));
     }
+
+    @Test
+    public void testBatchAnalysis()
+    {
+        SqlAnalysisInputBatchDto inputBatchDto = new SqlAnalysisInputBatchDto(
+                base64Encode(toJson(manifest)),
+                List.of("select * from customer",
+                        "select custkey, count(*) from customer group by 1",
+                        "with t1 as (select * from customer) select * from t1",
+                        "select * from orders where orderstatus = 'O' union select * from orders where orderstatus = 'F'"));
+
+        List<List<QueryAnalysisDto>> results = getSqlAnalysisBatch(inputBatchDto);
+        assertThat(results.size()).isEqualTo(4);
+        assertThat(results.get(0).size()).isEqualTo(1);
+        assertThat(results.get(1).size()).isEqualTo(1);
+        assertThat(results.get(2).size()).isEqualTo(2);
+        assertThat(results.get(3).size()).isEqualTo(2);
+    }
+
+    private String toJson(Manifest manifest)
+    {
+        return MANIFEST_JSON_CODEC.toJson(manifest);
+    }
+
+    private String base64Encode(String str)
+    {
+        return Base64.getEncoder().encodeToString(str.getBytes(UTF_8));
+    }
 }