Skip to content

Commit 744eb72

Browse files
committed
Convert S3HudiQueryRunner to builder
1 parent b42773b commit 744eb72

File tree

3 files changed

+81
-73
lines changed

3 files changed

+81
-73
lines changed

plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/S3HudiQueryRunner.java

Lines changed: 73 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@
1313
*/
1414
package io.trino.plugin.hudi;
1515

16-
import com.google.common.collect.ImmutableMap;
16+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
1717
import io.airlift.log.Logger;
1818
import io.airlift.log.Logging;
19-
import io.trino.Session;
2019
import io.trino.filesystem.Location;
20+
import io.trino.plugin.base.util.Closables;
2121
import io.trino.plugin.hive.containers.HiveMinioDataLake;
2222
import io.trino.plugin.hive.metastore.Database;
2323
import io.trino.plugin.hive.metastore.HiveMetastoreFactory;
@@ -28,81 +28,92 @@
2828
import io.trino.testing.QueryRunner;
2929
import io.trino.tpch.TpchTable;
3030

31+
import java.util.HashMap;
3132
import java.util.Map;
3233
import java.util.Optional;
3334

3435
import static io.trino.testing.TestingSession.testSessionBuilder;
3536
import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY;
3637
import static io.trino.testing.containers.Minio.MINIO_REGION;
3738
import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY;
39+
import static java.util.Objects.requireNonNull;
3840

41+
// TODO merge with HudiQueryRunner
3942
public final class S3HudiQueryRunner
4043
{
4144
private static final String TPCH_SCHEMA = "tpch";
4245

4346
private S3HudiQueryRunner() {}
4447

45-
// TODO convert to builder, merge with HudiQueryRunner
46-
public static QueryRunner create(
47-
Map<String, String> connectorProperties,
48-
HudiTablesInitializer dataLoader,
49-
HiveMinioDataLake hiveMinioDataLake)
50-
throws Exception
48+
public static Builder builder(HiveMinioDataLake hiveMinioDataLake)
5149
{
52-
return create(
53-
ImmutableMap.of(),
54-
connectorProperties,
55-
dataLoader,
56-
hiveMinioDataLake);
50+
return new Builder(hiveMinioDataLake)
51+
.addConnectorProperty("fs.hadoop.enabled", "false")
52+
.addConnectorProperty("fs.native-s3.enabled", "true")
53+
.addConnectorProperty("s3.aws-access-key", MINIO_ACCESS_KEY)
54+
.addConnectorProperty("s3.aws-secret-key", MINIO_SECRET_KEY)
55+
.addConnectorProperty("s3.region", MINIO_REGION)
56+
.addConnectorProperty("s3.endpoint", hiveMinioDataLake.getMinio().getMinioAddress())
57+
.addConnectorProperty("s3.path-style-access", "true");
5758
}
5859

59-
// TODO convert to builder, merge with HudiQueryRunner
60-
private static QueryRunner create(
61-
Map<String, String> coordinatorProperties,
62-
Map<String, String> connectorProperties,
63-
HudiTablesInitializer dataLoader,
64-
HiveMinioDataLake hiveMinioDataLake)
65-
throws Exception
60+
public static class Builder
61+
extends DistributedQueryRunner.Builder<Builder>
6662
{
67-
QueryRunner queryRunner = DistributedQueryRunner.builder(createSession())
68-
.setCoordinatorProperties(coordinatorProperties)
69-
.build();
70-
queryRunner.installPlugin(new TestingHudiPlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hudi_data")));
71-
queryRunner.createCatalog(
72-
"hudi",
73-
"hudi",
74-
ImmutableMap.<String, String>builder()
75-
.put("fs.hadoop.enabled", "false")
76-
.put("fs.native-s3.enabled", "true")
77-
.put("s3.aws-access-key", MINIO_ACCESS_KEY)
78-
.put("s3.aws-secret-key", MINIO_SECRET_KEY)
79-
.put("s3.region", MINIO_REGION)
80-
.put("s3.endpoint", hiveMinioDataLake.getMinio().getMinioAddress())
81-
.put("s3.path-style-access", "true")
82-
.putAll(connectorProperties)
83-
.buildOrThrow());
84-
85-
// Hudi connector does not support creating schema or any other write operations
86-
((HudiConnector) queryRunner.getCoordinator().getConnector("hudi")).getInjector()
87-
.getInstance(HiveMetastoreFactory.class)
88-
.createMetastore(Optional.empty())
89-
.createDatabase(Database.builder()
90-
.setDatabaseName(TPCH_SCHEMA)
91-
.setOwnerName(Optional.of("public"))
92-
.setOwnerType(Optional.of(PrincipalType.ROLE))
93-
.build());
94-
95-
dataLoader.initializeTables(queryRunner, Location.of("s3://" + hiveMinioDataLake.getBucketName() + "/"), TPCH_SCHEMA);
96-
97-
return queryRunner;
98-
}
63+
private final HiveMinioDataLake hiveMinioDataLake;
64+
private HudiTablesInitializer dataLoader;
65+
private final Map<String, String> connectorProperties = new HashMap<>();
9966

100-
private static Session createSession()
101-
{
102-
return testSessionBuilder()
103-
.setCatalog("hudi")
104-
.setSchema(TPCH_SCHEMA)
105-
.build();
67+
protected Builder(HiveMinioDataLake hiveMinioDataLake)
68+
{
69+
super(testSessionBuilder()
70+
.setCatalog("hudi")
71+
.setSchema(TPCH_SCHEMA)
72+
.build());
73+
this.hiveMinioDataLake = requireNonNull(hiveMinioDataLake, "hiveMinioDataLake is null");
74+
}
75+
76+
@CanIgnoreReturnValue
77+
public Builder setDataLoader(HudiTablesInitializer dataLoader)
78+
{
79+
this.dataLoader = dataLoader;
80+
return this;
81+
}
82+
83+
@CanIgnoreReturnValue
84+
public Builder addConnectorProperty(String key, String value)
85+
{
86+
this.connectorProperties.put(key, value);
87+
return this;
88+
}
89+
90+
@Override
91+
public DistributedQueryRunner build()
92+
throws Exception
93+
{
94+
DistributedQueryRunner queryRunner = super.build();
95+
try {
96+
queryRunner.installPlugin(new TestingHudiPlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hudi_data")));
97+
queryRunner.createCatalog("hudi", "hudi", connectorProperties);
98+
99+
// Hudi connector does not support creating schema or any other write operations
100+
((HudiConnector) queryRunner.getCoordinator().getConnector("hudi")).getInjector()
101+
.getInstance(HiveMetastoreFactory.class)
102+
.createMetastore(Optional.empty())
103+
.createDatabase(Database.builder()
104+
.setDatabaseName(TPCH_SCHEMA)
105+
.setOwnerName(Optional.of("public"))
106+
.setOwnerType(Optional.of(PrincipalType.ROLE))
107+
.build());
108+
109+
dataLoader.initializeTables(queryRunner, Location.of("s3://" + hiveMinioDataLake.getBucketName() + "/"), TPCH_SCHEMA);
110+
return queryRunner;
111+
}
112+
catch (Throwable e) {
113+
Closables.closeAllSuppress(e, queryRunner);
114+
throw e;
115+
}
116+
}
106117
}
107118

108119
public static void main(String[] args)
@@ -114,11 +125,10 @@ public static void main(String[] args)
114125
String bucketName = "test-bucket";
115126
HiveMinioDataLake hiveMinioDataLake = new HiveMinioDataLake(bucketName);
116127
hiveMinioDataLake.start();
117-
QueryRunner queryRunner = create(
118-
ImmutableMap.of("http-server.http.port", "8080"),
119-
ImmutableMap.of(),
120-
new TpchHudiTablesInitializer(TpchTable.getTables()),
121-
hiveMinioDataLake);
128+
QueryRunner queryRunner = builder(hiveMinioDataLake)
129+
.addCoordinatorProperty("http-server.http.port", "8080")
130+
.setDataLoader(new TpchHudiTablesInitializer(TpchTable.getTables()))
131+
.build();
122132

123133
log.info("======== SERVER STARTED ========");
124134
log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());

plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiCopyOnWriteMinioConnectorSmokeTest.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
*/
1414
package io.trino.plugin.hudi;
1515

16-
import com.google.common.collect.ImmutableMap;
1716
import io.trino.plugin.hive.containers.HiveMinioDataLake;
1817
import io.trino.plugin.hudi.testing.TpchHudiTablesInitializer;
1918
import io.trino.testing.QueryRunner;
@@ -34,9 +33,9 @@ protected QueryRunner createQueryRunner()
3433
hiveMinioDataLake.start();
3534
hiveMinioDataLake.getMinioClient().ensureBucketExists(bucketName);
3635

37-
return S3HudiQueryRunner.create(
38-
ImmutableMap.of("hudi.columns-to-hide", COLUMNS_TO_HIDE),
39-
new TpchHudiTablesInitializer(REQUIRED_TPCH_TABLES),
40-
hiveMinioDataLake);
36+
return S3HudiQueryRunner.builder(hiveMinioDataLake)
37+
.addConnectorProperty("hudi.columns-to-hide", COLUMNS_TO_HIDE)
38+
.setDataLoader(new TpchHudiTablesInitializer(REQUIRED_TPCH_TABLES))
39+
.build();
4140
}
4241
}

plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiMergeOnReadMinioConnectorSmokeTest.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
*/
1414
package io.trino.plugin.hudi;
1515

16-
import com.google.common.collect.ImmutableMap;
1716
import io.trino.plugin.hive.containers.HiveMinioDataLake;
1817
import io.trino.plugin.hudi.testing.TpchHudiTablesInitializer;
1918
import io.trino.testing.QueryRunner;
@@ -34,9 +33,9 @@ protected QueryRunner createQueryRunner()
3433
hiveMinioDataLake.start();
3534
hiveMinioDataLake.getMinioClient().ensureBucketExists(bucketName);
3635

37-
return S3HudiQueryRunner.create(
38-
ImmutableMap.of("hudi.columns-to-hide", COLUMNS_TO_HIDE),
39-
new TpchHudiTablesInitializer(REQUIRED_TPCH_TABLES),
40-
hiveMinioDataLake);
36+
return S3HudiQueryRunner.builder(hiveMinioDataLake)
37+
.setDataLoader(new TpchHudiTablesInitializer(REQUIRED_TPCH_TABLES))
38+
.addConnectorProperty("hudi.columns-to-hide", COLUMNS_TO_HIDE)
39+
.build();
4140
}
4241
}

0 commit comments

Comments
 (0)