🌐 AI搜索 & 代理 主页
Skip to content

Commit 281e7b1

Browse files
hashharlosipiuk
hashhar
authored and
losipiuk
committed
Add BigQuerySqlExecutor and refactor tests to use it
1 parent cc90d10 commit 281e7b1

File tree

2 files changed

+72
-82
lines changed

2 files changed

+72
-82
lines changed

‎plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/BigQueryQueryRunner.java‎

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616
import com.google.auth.oauth2.ServiceAccountCredentials;
1717
import com.google.cloud.bigquery.BigQuery;
1818
import com.google.cloud.bigquery.BigQueryOptions;
19+
import com.google.cloud.bigquery.QueryJobConfiguration;
1920
import com.google.common.collect.ImmutableMap;
2021
import io.airlift.log.Logger;
2122
import io.airlift.log.Logging;
2223
import io.trino.Session;
2324
import io.trino.plugin.tpch.TpchPlugin;
2425
import io.trino.testing.DistributedQueryRunner;
26+
import io.trino.testing.sql.SqlExecutor;
2527

2628
import java.io.ByteArrayInputStream;
2729
import java.io.IOException;
@@ -33,7 +35,7 @@
3335
import static io.airlift.testing.Closeables.closeAllSuppress;
3436
import static io.trino.testing.TestingSession.testSessionBuilder;
3537

36-
public class BigQueryQueryRunner
38+
public final class BigQueryQueryRunner
3739
{
3840
private static final String TPCH_SCHEMA = "tpch";
3941

@@ -65,20 +67,6 @@ public static DistributedQueryRunner createQueryRunner(Map<String, String> extra
6567
}
6668
}
6769

68-
public static BigQuery createBigQueryClient()
69-
{
70-
try {
71-
InputStream jsonKey = new ByteArrayInputStream(Base64.getDecoder().decode(System.getProperty("bigquery.credentials-key")));
72-
return BigQueryOptions.newBuilder()
73-
.setCredentials(ServiceAccountCredentials.fromStream(jsonKey))
74-
.build()
75-
.getService();
76-
}
77-
catch (IOException e) {
78-
throw new UncheckedIOException(e);
79-
}
80-
}
81-
8270
public static Session createSession()
8371
{
8472
return testSessionBuilder()
@@ -87,6 +75,43 @@ public static Session createSession()
8775
.build();
8876
}
8977

78+
public static class BigQuerySqlExecutor
79+
implements SqlExecutor
80+
{
81+
private final BigQuery bigQuery;
82+
83+
public BigQuerySqlExecutor()
84+
{
85+
this.bigQuery = createBigQueryClient();
86+
}
87+
88+
@Override
89+
public void execute(String sql)
90+
{
91+
try {
92+
bigQuery.query(QueryJobConfiguration.of(sql));
93+
}
94+
catch (InterruptedException e) {
95+
Thread.currentThread().interrupt();
96+
throw new RuntimeException(e);
97+
}
98+
}
99+
100+
private static BigQuery createBigQueryClient()
101+
{
102+
try {
103+
InputStream jsonKey = new ByteArrayInputStream(Base64.getDecoder().decode(System.getProperty("bigquery.credentials-key")));
104+
return BigQueryOptions.newBuilder()
105+
.setCredentials(ServiceAccountCredentials.fromStream(jsonKey))
106+
.build()
107+
.getService();
108+
}
109+
catch (IOException e) {
110+
throw new UncheckedIOException(e);
111+
}
112+
}
113+
}
114+
90115
public static void main(String[] args)
91116
throws Exception
92117
{

‎plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java‎

Lines changed: 32 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,13 @@
1313
*/
1414
package io.trino.plugin.bigquery;
1515

16-
import com.google.cloud.bigquery.BigQuery;
17-
import com.google.cloud.bigquery.Job;
18-
import com.google.cloud.bigquery.JobId;
19-
import com.google.cloud.bigquery.JobInfo;
20-
import com.google.cloud.bigquery.QueryJobConfiguration;
2116
import com.google.common.collect.ImmutableMap;
2217
import io.trino.testing.AbstractTestIntegrationSmokeTest;
2318
import io.trino.testing.MaterializedResult;
2419
import io.trino.testing.QueryRunner;
2520
import org.testng.annotations.Test;
2621

27-
import static io.trino.plugin.bigquery.BigQueryQueryRunner.createBigQueryClient;
22+
import static io.trino.plugin.bigquery.BigQueryQueryRunner.BigQuerySqlExecutor;
2823
import static io.trino.spi.type.VarcharType.VARCHAR;
2924
import static io.trino.testing.MaterializedResult.resultBuilder;
3025
import static io.trino.testing.assertions.Assert.assertEquals;
@@ -36,10 +31,13 @@
3631
public class TestBigQueryIntegrationSmokeTest
3732
extends AbstractTestIntegrationSmokeTest
3833
{
34+
private BigQuerySqlExecutor bigQuerySqlExecutor;
35+
3936
@Override
4037
protected QueryRunner createQueryRunner()
4138
throws Exception
4239
{
40+
this.bigQuerySqlExecutor = new BigQuerySqlExecutor();
4341
return BigQueryQueryRunner.createQueryRunner(ImmutableMap.of());
4442
}
4543

@@ -64,11 +62,9 @@ public void testDescribeTable()
6462
@Test(enabled = false)
6563
public void testSelectFromHourlyPartitionedTable()
6664
{
67-
BigQuery client = createBigQueryClient();
68-
69-
executeBigQuerySql(client, "DROP TABLE IF EXISTS test.hourly_partitioned");
70-
executeBigQuerySql(client, "CREATE TABLE test.hourly_partitioned (value INT64, ts TIMESTAMP) PARTITION BY TIMESTAMP_TRUNC(ts, HOUR)");
71-
executeBigQuerySql(client, "INSERT INTO test.hourly_partitioned (value, ts) VALUES (1000, '2018-01-01 10:00:00')");
65+
onBigQuery("DROP TABLE IF EXISTS test.hourly_partitioned");
66+
onBigQuery("CREATE TABLE test.hourly_partitioned (value INT64, ts TIMESTAMP) PARTITION BY TIMESTAMP_TRUNC(ts, HOUR)");
67+
onBigQuery("INSERT INTO test.hourly_partitioned (value, ts) VALUES (1000, '2018-01-01 10:00:00')");
7268

7369
MaterializedResult actualValues = computeActual("SELECT COUNT(1) FROM test.hourly_partitioned");
7470

@@ -78,11 +74,9 @@ public void testSelectFromHourlyPartitionedTable()
7874
@Test(enabled = false)
7975
public void testSelectFromYearlyPartitionedTable()
8076
{
81-
BigQuery client = createBigQueryClient();
82-
83-
executeBigQuerySql(client, "DROP TABLE IF EXISTS test.yearly_partitioned");
84-
executeBigQuerySql(client, "CREATE TABLE test.yearly_partitioned (value INT64, ts TIMESTAMP) PARTITION BY TIMESTAMP_TRUNC(ts, YEAR)");
85-
executeBigQuerySql(client, "INSERT INTO test.yearly_partitioned (value, ts) VALUES (1000, '2018-01-01 10:00:00')");
77+
onBigQuery("DROP TABLE IF EXISTS test.yearly_partitioned");
78+
onBigQuery("CREATE TABLE test.yearly_partitioned (value INT64, ts TIMESTAMP) PARTITION BY TIMESTAMP_TRUNC(ts, YEAR)");
79+
onBigQuery("INSERT INTO test.yearly_partitioned (value, ts) VALUES (1000, '2018-01-01 10:00:00')");
8680

8781
MaterializedResult actualValues = computeActual("SELECT COUNT(1) FROM test.yearly_partitioned");
8882

@@ -92,13 +86,11 @@ public void testSelectFromYearlyPartitionedTable()
9286
@Test(description = "regression test for https://github.com/trinodb/trino/issues/5618")
9387
public void testPredicatePushdownPrunnedColumns()
9488
{
95-
BigQuery client = createBigQueryClient();
96-
9789
String tableName = "test.predicate_pushdown_prunned_columns";
9890

99-
executeBigQuerySql(client, "DROP TABLE IF EXISTS " + tableName);
100-
executeBigQuerySql(client, "CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
101-
executeBigQuerySql(client, "INSERT INTO " + tableName + " VALUES (1,2,3)");
91+
onBigQuery("DROP TABLE IF EXISTS " + tableName);
92+
onBigQuery("CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
93+
onBigQuery("INSERT INTO " + tableName + " VALUES (1,2,3)");
10294

10395
assertQuery(
10496
"SELECT 1 FROM " + tableName + " WHERE " +
@@ -110,16 +102,14 @@ public void testPredicatePushdownPrunnedColumns()
110102
@Test(description = "regression test for https://github.com/trinodb/trino/issues/5635")
111103
public void testCountAggregationView()
112104
{
113-
BigQuery client = createBigQueryClient();
114-
115105
String tableName = "test.count_aggregation_table";
116106
String viewName = "test.count_aggregation_view";
117107

118-
executeBigQuerySql(client, "DROP TABLE IF EXISTS " + tableName);
119-
executeBigQuerySql(client, "DROP VIEW IF EXISTS " + viewName);
120-
executeBigQuerySql(client, "CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
121-
executeBigQuerySql(client, "INSERT INTO " + tableName + " VALUES (1, 2, 3), (4, 5, 6)");
122-
executeBigQuerySql(client, "CREATE VIEW " + viewName + " AS SELECT * FROM " + tableName);
108+
onBigQuery("DROP TABLE IF EXISTS " + tableName);
109+
onBigQuery("DROP VIEW IF EXISTS " + viewName);
110+
onBigQuery("CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
111+
onBigQuery("INSERT INTO " + tableName + " VALUES (1, 2, 3), (4, 5, 6)");
112+
onBigQuery("CREATE VIEW " + viewName + " AS SELECT * FROM " + tableName);
123113

124114
assertQuery(
125115
"SELECT count(*) FROM " + viewName,
@@ -140,32 +130,28 @@ public void testCountAggregationView()
140130
@Test
141131
public void testRepeatCountAggregationView()
142132
{
143-
BigQuery client = createBigQueryClient();
144-
145133
String viewName = "test.repeat_count_aggregation_view_" + randomTableSuffix();
146134

147-
executeBigQuerySql(client, "DROP VIEW IF EXISTS " + viewName);
148-
executeBigQuerySql(client, "CREATE VIEW " + viewName + " AS SELECT 1 AS col1");
135+
onBigQuery("DROP VIEW IF EXISTS " + viewName);
136+
onBigQuery("CREATE VIEW " + viewName + " AS SELECT 1 AS col1");
149137

150138
assertQuery("SELECT count(*) FROM " + viewName, "VALUES (1)");
151139
assertQuery("SELECT count(*) FROM " + viewName, "VALUES (1)");
152140

153-
executeBigQuerySql(client, "DROP VIEW " + viewName);
141+
onBigQuery("DROP VIEW " + viewName);
154142
}
155143

156144
@Test
157145
public void testViewDefinitionSystemTable()
158146
{
159-
BigQuery client = createBigQueryClient();
160-
161147
String schemaName = "test";
162148
String tableName = "views_system_table_base_" + randomTableSuffix();
163149
String viewName = "views_system_table_view_" + randomTableSuffix();
164150

165-
executeBigQuerySql(client, format("DROP TABLE IF EXISTS %s.%s", schemaName, tableName));
166-
executeBigQuerySql(client, format("DROP VIEW IF EXISTS %s.%s", schemaName, viewName));
167-
executeBigQuerySql(client, format("CREATE TABLE %s.%s (a INT64, b INT64, c INT64)", schemaName, tableName));
168-
executeBigQuerySql(client, format("CREATE VIEW %s.%s AS SELECT * FROM %s.%s", schemaName, viewName, schemaName, tableName));
151+
onBigQuery(format("DROP TABLE IF EXISTS %s.%s", schemaName, tableName));
152+
onBigQuery(format("DROP VIEW IF EXISTS %s.%s", schemaName, viewName));
153+
onBigQuery(format("CREATE TABLE %s.%s (a INT64, b INT64, c INT64)", schemaName, tableName));
154+
onBigQuery(format("CREATE VIEW %s.%s AS SELECT * FROM %s.%s", schemaName, viewName, schemaName, tableName));
169155

170156
assertEquals(
171157
computeScalar(format("SELECT * FROM %s.\"%s$view_definition\"", schemaName, viewName)),
@@ -175,34 +161,8 @@ public void testViewDefinitionSystemTable()
175161
format("SELECT * FROM %s.\"%s$view_definition\"", schemaName, tableName),
176162
format("Table '%s.%s\\$view_definition' not found", schemaName, tableName));
177163

178-
executeBigQuerySql(client, format("DROP TABLE %s.%s", schemaName, tableName));
179-
executeBigQuerySql(client, format("DROP VIEW %s.%s", schemaName, viewName));
180-
}
181-
182-
private static void executeBigQuerySql(BigQuery bigquery, String query)
183-
{
184-
QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query)
185-
.setUseLegacySql(false)
186-
.build();
187-
188-
JobId jobId = JobId.of();
189-
Job queryJob = bigquery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());
190-
191-
try {
192-
queryJob = queryJob.waitFor();
193-
194-
if (queryJob == null) {
195-
throw new RuntimeException(format("Job with uuid %s does not longer exists", jobId.getJob()));
196-
}
197-
198-
if (queryJob.getStatus().getError() != null) {
199-
throw new RuntimeException(format("Query '%s' failed: %s", query, queryJob.getStatus().getError()));
200-
}
201-
}
202-
catch (InterruptedException e) {
203-
Thread.currentThread().interrupt();
204-
throw new RuntimeException(e);
205-
}
164+
onBigQuery(format("DROP TABLE %s.%s", schemaName, tableName));
165+
onBigQuery(format("DROP VIEW %s.%s", schemaName, viewName));
206166
}
207167

208168
@Override
@@ -221,4 +181,9 @@ public void testShowCreateTable()
221181
" comment varchar NOT NULL\n" +
222182
")");
223183
}
184+
185+
private void onBigQuery(String sql)
186+
{
187+
bigQuerySqlExecutor.execute(sql);
188+
}
224189
}

0 commit comments

Comments
 (0)