Provide API for SQL Dry-Run and Dry-Plan

wren-main/src/main/java/io/wren/main/PreviewService.java

@@ -31,6 +31,7 @@
 import com.google.common.collect.Streams;
 import com.google.inject.Inject;
 import io.wren.base.AnalyzedMDL;
+import io.wren.base.Column;
 import io.wren.base.ConnectorRecordIterator;
 import io.wren.base.SessionContext;
 import io.wren.base.WrenMDL;
@@ -80,4 +81,34 @@ public CompletableFuture<QueryResultDto> preview(WrenMDL mdl, String sql, long l
             }
         });
     }
+
+    public CompletableFuture<String> dryPlan(WrenMDL mdl, String sql, boolean isModelingOnly)
+    {
+        return CompletableFuture.supplyAsync(() -> {
+            SessionContext sessionContext = SessionContext.builder()
+                    .setCatalog(mdl.getCatalog())
+                    .setSchema(mdl.getSchema())
+                    .build();
+
+            String planned = WrenPlanner.rewrite(sql, sessionContext, new AnalyzedMDL(mdl, null));
+            if (isModelingOnly) {
+                return planned;
+            }
+            return sqlConverter.convert(planned, sessionContext);
+        });
+    }
+
+    public CompletableFuture<List<Column>> dryRun(WrenMDL mdl, String sql)
+    {
+        return CompletableFuture.supplyAsync(() -> {
+            SessionContext sessionContext = SessionContext.builder()
+                    .setCatalog(mdl.getCatalog())
+                    .setSchema(mdl.getSchema())
+                    .build();
+
+            String planned = WrenPlanner.rewrite(sql, sessionContext, new AnalyzedMDL(mdl, null));
+            String converted = sqlConverter.convert(planned, sessionContext);
+            return metadata.describeQuery(converted, List.of());
+        });
+    }
 }

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

@@ -20,6 +20,7 @@
 import io.wren.main.WrenManager;
 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.PreviewDto;
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.GET;
@@ -107,4 +108,42 @@ public void preview(
                         Optional.ofNullable(previewDto.getLimit()).orElse(100L))
                 .whenComplete(bindAsyncResponse(asyncResponse));
     }
+
+    @GET
+    @Path("/dry-plan")
+    @Consumes(APPLICATION_JSON)
+    @Produces(APPLICATION_JSON)
+    public void dryPlan(
+            DryPlanDto dryPlanDto,
+            @Suspended AsyncResponse asyncResponse)
+    {
+        WrenMDL mdl;
+        if (dryPlanDto.getManifest() == null) {
+            mdl = wrenManager.getAnalyzedMDL().getWrenMDL();
+        }
+        else {
+            mdl = WrenMDL.fromManifest(dryPlanDto.getManifest());
+        }
+        previewService.dryPlan(mdl, dryPlanDto.getSql(), dryPlanDto.isModelingOnly())
+                .whenComplete(bindAsyncResponse(asyncResponse));
+    }
+
+    @GET
+    @Path("/dry-run")
+    @Consumes(APPLICATION_JSON)
+    @Produces(APPLICATION_JSON)
+    public void dryRun(
+            PreviewDto previewDto,
+            @Suspended AsyncResponse asyncResponse)
+    {
+        WrenMDL mdl;
+        if (previewDto.getManifest() == null) {
+            mdl = wrenManager.getAnalyzedMDL().getWrenMDL();
+        }
+        else {
+            mdl = WrenMDL.fromManifest(previewDto.getManifest());
+        }
+        previewService.dryRun(mdl, previewDto.getSql())
+                .whenComplete(bindAsyncResponse(asyncResponse));
+    }
 }

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

@@ -0,0 +1,55 @@
+/*
+ * 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 io.wren.base.dto.Manifest;
+
+public class DryPlanDto
+{
+    private final Manifest manifest;
+    private final String sql;
+    private final boolean isModelingOnly;
+
+    @JsonCreator
+    public DryPlanDto(
+            @JsonProperty("manifest") Manifest manifest,
+            @JsonProperty("sql") String sql,
+            @JsonProperty("modelingOnly") boolean modelingOnly)
+    {
+        this.manifest = manifest;
+        this.sql = sql;
+        this.isModelingOnly = modelingOnly;
+    }
+
+    @JsonProperty
+    public Manifest getManifest()
+    {
+        return manifest;
+    }
+
+    @JsonProperty
+    public String getSql()
+    {
+        return sql;
+    }
+
+    @JsonProperty
+    public boolean isModelingOnly()
+    {
+        return isModelingOnly;
+    }
+}

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

@@ -26,11 +26,13 @@
 import io.airlift.units.Duration;
 import io.wren.base.CatalogSchemaTableName;
 import io.wren.base.config.ConfigManager;
+import io.wren.base.dto.Column;
 import io.wren.base.dto.Manifest;
 import io.wren.cache.TaskInfo;
 import io.wren.main.web.dto.CheckOutputDto;
 import io.wren.main.web.dto.ColumnLineageInputDto;
 import io.wren.main.web.dto.DeployInputDto;
+import io.wren.main.web.dto.DryPlanDto;
 import io.wren.main.web.dto.ErrorMessageDto;
 import io.wren.main.web.dto.LineageResult;
 import io.wren.main.web.dto.PreviewDto;
@@ -85,6 +87,8 @@ public abstract class RequireWrenServer
     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);
     private static final JsonCodec<QueryResultDto> QUERY_RESULT_DTO_CODEC = jsonCodec(QueryResultDto.class);
+    private static final JsonCodec<List<Column>> COLUMN_LIST_CODEC = listJsonCodec(Column.class);
+    private static final JsonCodec<DryPlanDto> DRY_PLAN_DTO_CODEC = jsonCodec(DryPlanDto.class);
 
     public RequireWrenServer() {}
 
@@ -170,6 +174,36 @@ protected QueryResultDto preview(PreviewDto previewDto)
         return QUERY_RESULT_DTO_CODEC.fromJson(response.getBody());
     }
 
+    protected List<Column> dryRun(PreviewDto previewDto)
+    {
+        Request request = prepareGet()
+                .setUri(server().getHttpServerBasedUrl().resolve("/v1/mdl/dry-run"))
+                .setHeader(CONTENT_TYPE, "application/json")
+                .setBodyGenerator(jsonBodyGenerator(PREVIEW_DTO_CODEC, previewDto))
+                .build();
+
+        StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
+        if (response.getStatusCode() != 200) {
+            getWebApplicationException(response);
+        }
+        return COLUMN_LIST_CODEC.fromJson(response.getBody());
+    }
+
+    protected String dryPlan(DryPlanDto dryPlanDto)
+    {
+        Request request = prepareGet()
+                .setUri(server().getHttpServerBasedUrl().resolve("/v1/mdl/dry-plan"))
+                .setHeader(CONTENT_TYPE, "application/json")
+                .setBodyGenerator(jsonBodyGenerator(DRY_PLAN_DTO_CODEC, dryPlanDto))
+                .build();
+
+        StringResponseHandler.StringResponse response = executeHttpRequest(request, createStringResponseHandler());
+        if (response.getStatusCode() != 200) {
+            getWebApplicationException(response);
+        }
+        return response.getBody();
+    }
+
     protected void deployMDL(DeployInputDto dto)
     {
         Request request = preparePost()

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

@@ -15,10 +15,12 @@
 package io.wren.testing;
 
 import com.google.common.collect.ImmutableMap;
+import io.wren.base.dto.Column;
 import io.wren.base.dto.Manifest;
 import io.wren.base.type.BigIntType;
 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.PreviewDto;
 import io.wren.main.web.dto.QueryResultDto;
 import org.testng.annotations.Test;
@@ -141,4 +143,67 @@ public void testPreview()
         assertWebApplicationException(() -> preview(new PreviewDto(previewManifest, "select orderkey from Orders limit 100", null)))
                 .hasErrorMessageMatches(".*Table \"Orders\" must be qualified with a dataset.*");
     }
+
+    @Test
+    public void testDryRunAndDryPlan()
+    {
+        Manifest previewManifest = Manifest.builder()
+                .setCatalog("canner-cml")
+                .setSchema("tpch_tiny")
+                .setModels(List.of(
+                        model("Customer", "SELECT * FROM \"canner-cml\".tpch_tiny.customer",
+                                List.of(column("custkey", "integer", null, false, "c_custkey")))))
+                .build();
+
+        PreviewDto testDefaultDto1 = new PreviewDto(previewManifest, "select custkey from Customer limit 200", null);
+        List<Column> dryRun = dryRun(testDefaultDto1);
+        assertThat(dryRun.size()).isEqualTo(1);
+        assertThat(dryRun.get(0).getName()).isEqualTo("custkey");
+
+        DryPlanDto dryPlanDto = new DryPlanDto(previewManifest, "select custkey from Customer limit 200", false);
+        String dryPlan = dryPlan(dryPlanDto);
+        assertThat(dryPlan).isEqualTo("""
+                WITH
+                  `Customer` AS (
+                   SELECT `Customer`.`custkey` `custkey`
+                   FROM
+                     (
+                      SELECT c_custkey `custkey`
+                      FROM
+                        (
+                         SELECT *
+                         FROM
+                           `canner-cml`.tpch_tiny.customer
+                      )  `Customer`
+                   )  `Customer`
+                )\s
+                SELECT custkey
+                FROM
+                  Customer
+                LIMIT 200
+                """);
+
+        dryPlanDto = new DryPlanDto(previewManifest, "select custkey from Customer limit 200", true);
+        dryPlan = dryPlan(dryPlanDto);
+        assertThat(dryPlan).isEqualTo("""
+                WITH
+                  "Customer" AS (
+                   SELECT "Customer"."custkey" "custkey"
+                   FROM
+                     (
+                      SELECT c_custkey "custkey"
+                      FROM
+                        (
+                         SELECT *
+                         FROM
+                           "canner-cml".tpch_tiny.customer
+                      )  "Customer"
+                   )  "Customer"
+                )\s
+                SELECT custkey
+                FROM
+                  Customer
+                LIMIT 200
+                """);
+    }
 }