Skip to content

Commit 453094d

Browse files
goldmedalgoldmedal
authored
Support dynamic query mode in preview api (#532)
* consider dynamic query mode in preview * add test and refactor to testing based DuckDB
1 parent 9d0bcdd commit 453094d

File tree

3 files changed

+151
-42
lines changed

3 files changed

+151
-42
lines changed

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

+12-1
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535
import io.wren.base.ConnectorRecordIterator;
3636
import io.wren.base.SessionContext;
3737
import io.wren.base.WrenMDL;
38+
import io.wren.base.config.ConfigManager;
39+
import io.wren.base.config.WrenConfig;
3840
import io.wren.base.sql.SqlConverter;
3941
import io.wren.base.sqlrewrite.WrenPlanner;
4042
import io.wren.main.metadata.Metadata;
@@ -51,22 +53,27 @@ public class PreviewService
5153
private final Metadata metadata;
5254

5355
private final SqlConverter sqlConverter;
56+
private final ConfigManager configManager;
5457

5558
@Inject
5659
public PreviewService(
5760
Metadata metadata,
58-
SqlConverter sqlConverter)
61+
SqlConverter sqlConverter,
62+
ConfigManager configManager)
5963
{
6064
this.metadata = requireNonNull(metadata, "metadata is null");
6165
this.sqlConverter = requireNonNull(sqlConverter, "sqlConverter is null");
66+
this.configManager = requireNonNull(configManager, "configManager is null");
6267
}
6368

6469
public CompletableFuture<QueryResultDto> preview(WrenMDL mdl, String sql, long limit)
6570
{
6671
return CompletableFuture.supplyAsync(() -> {
72+
WrenConfig config = configManager.getConfig(WrenConfig.class);
6773
SessionContext sessionContext = SessionContext.builder()
6874
.setCatalog(mdl.getCatalog())
6975
.setSchema(mdl.getSchema())
76+
.setEnableDynamic(config.getEnableDynamicFields())
7077
.build();
7178

7279
String planned = WrenPlanner.rewrite(sql, sessionContext, new AnalyzedMDL(mdl, null));
@@ -85,9 +92,11 @@ public CompletableFuture<QueryResultDto> preview(WrenMDL mdl, String sql, long l
8592
public CompletableFuture<String> dryPlan(WrenMDL mdl, String sql, boolean isModelingOnly)
8693
{
8794
return CompletableFuture.supplyAsync(() -> {
95+
WrenConfig config = configManager.getConfig(WrenConfig.class);
8896
SessionContext sessionContext = SessionContext.builder()
8997
.setCatalog(mdl.getCatalog())
9098
.setSchema(mdl.getSchema())
99+
.setEnableDynamic(config.getEnableDynamicFields())
91100
.build();
92101

93102
String planned = WrenPlanner.rewrite(sql, sessionContext, new AnalyzedMDL(mdl, null));
@@ -101,9 +110,11 @@ public CompletableFuture<String> dryPlan(WrenMDL mdl, String sql, boolean isMode
101110
public CompletableFuture<List<Column>> dryRun(WrenMDL mdl, String sql)
102111
{
103112
return CompletableFuture.supplyAsync(() -> {
113+
WrenConfig config = configManager.getConfig(WrenConfig.class);
104114
SessionContext sessionContext = SessionContext.builder()
105115
.setCatalog(mdl.getCatalog())
106116
.setSchema(mdl.getSchema())
117+
.setEnableDynamic(config.getEnableDynamicFields())
107118
.build();
108119

109120
String planned = WrenPlanner.rewrite(sql, sessionContext, new AnalyzedMDL(mdl, null));

‎wren-main/src/main/java/io/wren/main/WrenManager.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ private void archiveWrenMDL(WrenMDL oldWrenMDL)
137137
}
138138
}
139139
Files.copy(wrenMDLFile.toPath(),
140-
archived.toPath().resolve(wrenMDLFile.getName() + "." + LocalDateTime.now().format(DateTimeFormatter.ofPattern("uuuuMMddHHmmss"))));
140+
archived.toPath().resolve(wrenMDLFile.getName() + "." + LocalDateTime.now().format(DateTimeFormatter.ofPattern("uuuuMMddHHmmssnnnn"))));
141141
}
142142

143143
public boolean checkStatus()

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

+138-40
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,13 @@
1515
package io.wren.testing;
1616

1717
import com.google.common.collect.ImmutableMap;
18+
import com.google.common.io.Resources;
19+
import com.google.inject.Key;
1820
import io.wren.base.dto.Column;
21+
import io.wren.base.dto.JoinType;
1922
import io.wren.base.dto.Manifest;
20-
import io.wren.base.type.BigIntType;
23+
import io.wren.base.type.IntegerType;
24+
import io.wren.main.connector.duckdb.DuckDBMetadata;
2125
import io.wren.main.web.dto.CheckOutputDto;
2226
import io.wren.main.web.dto.DeployInputDto;
2327
import io.wren.main.web.dto.DryPlanDto;
@@ -30,14 +34,18 @@
3034
import java.nio.file.Path;
3135
import java.util.List;
3236

33-
import static io.wren.base.Utils.randomIntString;
3437
import static io.wren.base.client.duckdb.FileUtil.ARCHIVED;
38+
import static io.wren.base.config.WrenConfig.DataSourceType.DUCKDB;
39+
import static io.wren.base.config.WrenConfig.WREN_DATASOURCE_TYPE;
40+
import static io.wren.base.config.WrenConfig.WREN_DIRECTORY;
41+
import static io.wren.base.config.WrenConfig.WREN_ENABLE_DYNAMIC_FIELDS;
42+
import static io.wren.base.dto.Column.caluclatedColumn;
3543
import static io.wren.base.dto.Column.column;
3644
import static io.wren.base.dto.Manifest.MANIFEST_JSON_CODEC;
3745
import static io.wren.base.dto.Model.model;
46+
import static io.wren.base.dto.Relationship.relationship;
3847
import static io.wren.testing.WebApplicationExceptionAssert.assertWebApplicationException;
39-
import static java.lang.String.format;
40-
import static java.lang.System.getenv;
48+
import static java.nio.charset.StandardCharsets.UTF_8;
4149
import static java.util.Objects.requireNonNull;
4250
import static org.assertj.core.api.Assertions.assertThat;
4351
import static org.assertj.core.api.Assertions.assertThatNoException;
@@ -49,22 +57,23 @@ public class TestMDLResource
4957
private Path mdlDir;
5058
private Manifest initial = Manifest.builder()
5159
.setCatalog("wrenai")
52-
.setSchema("tpch_tiny")
60+
.setSchema("tpch")
5361
.setModels(List.of(
54-
model("Orders", "SELECT * FROM wrenai.tpch_tiny.orders", List.of(column("orderkey", "integer", null, false, "o_orderkey")))))
62+
model("Orders", "SELECT * FROM tpch.orders", List.of(column("orderkey", "integer", null, false, "o_orderkey")))))
5563
.build();
5664

5765
private Manifest updated = Manifest.builder()
5866
.setCatalog("wrenai")
59-
.setSchema("tpch_tiny")
67+
.setSchema("tpch")
6068
.setModels(List.of(
61-
model("Orders", "SELECT * FROM wrenai.tpch_tiny.orders",
69+
model("Orders", "SELECT * FROM tpch.orders",
6270
List.of(column("orderkey", "integer", null, false, "o_orderkey"),
6371
column("custkey", "integer", null, false, "o_custkey")))))
6472
.build();
6573

6674
@Override
6775
protected TestingWrenServer createWrenServer()
76+
throws Exception
6877
{
6978
try {
7079
mdlDir = Files.createTempDirectory("wrenmdls");
@@ -76,16 +85,26 @@ protected TestingWrenServer createWrenServer()
7685
}
7786

7887
ImmutableMap.Builder<String, String> properties = ImmutableMap.<String, String>builder()
79-
.put("bigquery.project-id", getenv("TEST_BIG_QUERY_PROJECT_ID"))
80-
.put("bigquery.location", "asia-east1")
81-
.put("bigquery.credentials-key", getenv("TEST_BIG_QUERY_CREDENTIALS_BASE64_JSON"))
82-
.put("bigquery.metadata.schema.prefix", format("test_%s_", randomIntString()))
83-
.put("wren.directory", mdlDir.toAbsolutePath().toString())
84-
.put("wren.datasource.type", "bigquery");
85-
86-
return TestingWrenServer.builder()
88+
.put(WREN_DIRECTORY, mdlDir.toAbsolutePath().toString())
89+
.put(WREN_DATASOURCE_TYPE, DUCKDB.name())
90+
.put(WREN_ENABLE_DYNAMIC_FIELDS, "true");
91+
92+
TestingWrenServer testing = TestingWrenServer.builder()
8793
.setRequiredConfigs(properties.build())
8894
.build();
95+
initDuckDB(testing);
96+
return testing;
97+
}
98+
99+
protected void initDuckDB(TestingWrenServer wrenServer)
100+
throws Exception
101+
{
102+
ClassLoader classLoader = getClass().getClassLoader();
103+
String initSQL = Resources.toString(requireNonNull(classLoader.getResource("duckdb/init.sql")).toURI().toURL(), UTF_8);
104+
initSQL = initSQL.replaceAll("basePath", requireNonNull(classLoader.getResource("tpch/data")).getPath());
105+
DuckDBMetadata metadata = wrenServer.getInstance(Key.get(DuckDBMetadata.class));
106+
metadata.setInitSQL(initSQL);
107+
metadata.reload();
89108
}
90109

91110
@Test
@@ -117,9 +136,9 @@ public void testPreview()
117136
{
118137
Manifest previewManifest = Manifest.builder()
119138
.setCatalog("wrenai")
120-
.setSchema("tpch_tiny")
139+
.setSchema("tpch")
121140
.setModels(List.of(
122-
model("Customer", "SELECT * FROM wrenai.tpch_tiny.customer",
141+
model("Customer", "SELECT * FROM tpch.customer",
123142
List.of(column("custkey", "integer", null, false, "c_custkey")))))
124143
.build();
125144

@@ -128,7 +147,7 @@ public void testPreview()
128147
assertThat(testDefault.getData().size()).isEqualTo(100);
129148
assertThat(testDefault.getColumns().size()).isEqualTo(1);
130149
assertThat(testDefault.getColumns().get(0).getName()).isEqualTo("custkey");
131-
assertThat(testDefault.getColumns().get(0).getType()).isEqualTo(BigIntType.BIGINT);
150+
assertThat(testDefault.getColumns().get(0).getType()).isEqualTo(IntegerType.INTEGER);
132151

133152
PreviewDto testDefaultDto1 = new PreviewDto(previewManifest, "select custkey from Customer limit 200", null);
134153
QueryResultDto preview1 = preview(testDefaultDto1);
@@ -141,68 +160,147 @@ public void testPreview()
141160
assertThat(preview2.getColumns().size()).isEqualTo(1);
142161

143162
assertWebApplicationException(() -> preview(new PreviewDto(previewManifest, "select orderkey from Orders limit 100", null)))
144-
.hasErrorMessageMatches(".*Table \"Orders\" must be qualified with a dataset.*");
163+
.hasErrorMessageMatches(".*Orders does not exist.*\n.*\n.*\n.*");
145164
}
146165

147166
@Test
148167
public void testDryRunAndDryPlan()
149168
{
150169
Manifest previewManifest = Manifest.builder()
151170
.setCatalog("wrenai")
152-
.setSchema("tpch_tiny")
171+
.setSchema("tpch")
153172
.setModels(List.of(
154-
model("Customer", "SELECT * FROM \"wrenai\".tpch_tiny.customer",
155-
List.of(column("custkey", "integer", null, false, "c_custkey")))))
173+
model("Customer", "SELECT * FROM tpch.customer",
174+
List.of(column("custkey", "integer", null, false, "c_custkey"),
175+
column("name", "varchar", null, false, "c_name"))),
176+
model("Orders", "SELECT * FROM tpch.orders",
177+
List.of(column("orderkey", "integer", null, false, "o_orderkey"),
178+
column("custkey", "integer", null, false, "o_custkey"),
179+
column("customer", "Customer", "CustomerOrders", false),
180+
caluclatedColumn("customer_name", "varchar", "customer.name")),
181+
"orderkey")))
182+
.setRelationships(List.of(relationship("CustomerOrders", List.of("Customer", "Orders"), JoinType.ONE_TO_MANY, "Customer.custkey = Orders.custkey")))
156183
.build();
157184

158-
PreviewDto testDefaultDto1 = new PreviewDto(previewManifest, "select custkey from Customer limit 200", null);
185+
PreviewDto testDefaultDto1 = new PreviewDto(previewManifest, "select orderkey from Orders limit 200", null);
159186
List<Column> dryRun = dryRun(testDefaultDto1);
160187
assertThat(dryRun.size()).isEqualTo(1);
161-
assertThat(dryRun.get(0).getName()).isEqualTo("custkey");
188+
assertThat(dryRun.get(0).getName()).isEqualTo("orderkey");
162189

163-
DryPlanDto dryPlanDto = new DryPlanDto(previewManifest, "select custkey from Customer limit 200", false);
190+
DryPlanDto dryPlanDto = new DryPlanDto(previewManifest, "select orderkey from Orders limit 200", false);
164191
String dryPlan = dryPlan(dryPlanDto);
165192
assertThat(dryPlan).isEqualTo("""
166193
WITH
167-
`Customer` AS (
168-
SELECT `Customer`.`custkey` `custkey`
194+
"Orders" AS (
195+
SELECT
196+
"Orders"."orderkey" "orderkey"
197+
, "Orders"."custkey" "custkey"
198+
FROM
199+
(
200+
SELECT
201+
o_orderkey "orderkey"
202+
, o_custkey "custkey"
203+
FROM
204+
(
205+
SELECT *
206+
FROM
207+
tpch.orders
208+
) "Orders"
209+
) "Orders"
210+
)\s
211+
SELECT orderkey
212+
FROM
213+
Orders
214+
LIMIT 200
215+
""");
216+
217+
dryPlanDto = new DryPlanDto(previewManifest, "select orderkey from Orders limit 200", true);
218+
dryPlan = dryPlan(dryPlanDto);
219+
assertThat(dryPlan).isEqualTo("""
220+
WITH
221+
"Orders" AS (
222+
SELECT
223+
"Orders"."orderkey" "orderkey"
224+
, "Orders"."custkey" "custkey"
169225
FROM
170226
(
171-
SELECT c_custkey `custkey`
227+
SELECT
228+
o_orderkey "orderkey"
229+
, o_custkey "custkey"
172230
FROM
173231
(
174232
SELECT *
175233
FROM
176-
`wrenai`.tpch_tiny.customer
177-
) `Customer`
178-
) `Customer`
234+
tpch.orders
235+
) "Orders"
236+
) "Orders"
179237
)\s
180-
SELECT custkey
238+
SELECT orderkey
181239
FROM
182-
Customer
240+
Orders
183241
LIMIT 200
184242
""");
185243

186-
dryPlanDto = new DryPlanDto(previewManifest, "select custkey from Customer limit 200", true);
244+
dryPlanDto = new DryPlanDto(previewManifest, "select customer_name from Orders limit 200", false);
187245
dryPlan = dryPlan(dryPlanDto);
188246
assertThat(dryPlan).isEqualTo("""
189247
WITH
190248
"Customer" AS (
191-
SELECT "Customer"."custkey" "custkey"
249+
SELECT
250+
"Customer"."custkey" "custkey"
251+
, "Customer"."name" "name"
192252
FROM
193253
(
194-
SELECT c_custkey "custkey"
254+
SELECT
255+
c_custkey "custkey"
256+
, c_name "name"
195257
FROM
196258
(
197259
SELECT *
198260
FROM
199-
"wrenai".tpch_tiny.customer
261+
tpch.customer
200262
) "Customer"
201263
) "Customer"
202264
)\s
203-
SELECT custkey
265+
, "Orders" AS (
266+
SELECT
267+
"Orders"."orderkey" "orderkey"
268+
, "Orders"."custkey" "custkey"
269+
, "Orders_relationsub"."customer_name" "customer_name"
270+
FROM
271+
((
272+
SELECT
273+
o_orderkey "orderkey"
274+
, o_custkey "custkey"
275+
FROM
276+
(
277+
SELECT *
278+
FROM
279+
tpch.orders
280+
) "Orders"
281+
) "Orders"
282+
LEFT JOIN (
283+
SELECT
284+
"Orders"."orderkey"
285+
, "Customer"."name" "customer_name"
286+
FROM
287+
((
288+
SELECT
289+
o_orderkey "orderkey"
290+
, o_custkey "custkey"
291+
FROM
292+
(
293+
SELECT *
294+
FROM
295+
tpch.orders
296+
) "Orders"
297+
) "Orders"
298+
LEFT JOIN "Customer" ON ("Customer"."custkey" = "Orders"."custkey"))
299+
) "Orders_relationsub" ON ("Orders"."orderkey" = "Orders_relationsub"."orderkey"))
300+
)\s
301+
SELECT customer_name
204302
FROM
205-
Customer
303+
Orders
206304
LIMIT 200
207305
""");
208306
}

0 commit comments

Comments
 (0)