Add BigQuery smoke test for HOUR-ly partitioned table

trinodb · Sep 1, 2020 · ae6c21b · ae6c21b
1 parent eb59569
commit ae6c21b
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 0 deletions.
diff --git a/presto-bigquery/src/test/java/io/prestosql/plugin/bigquery/BigQueryQueryRunner.java b/presto-bigquery/src/test/java/io/prestosql/plugin/bigquery/BigQueryQueryRunner.java
@@ -13,13 +13,21 @@
  */
 package io.prestosql.plugin.bigquery;
 
+import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryOptions;
 import com.google.common.collect.ImmutableMap;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
 import io.prestosql.Session;
 import io.prestosql.plugin.tpch.TpchPlugin;
 import io.prestosql.testing.DistributedQueryRunner;
 
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.util.Base64;
 import java.util.Map;
 
 import static io.airlift.testing.Closeables.closeAllSuppress;
@@ -54,6 +62,20 @@ public static DistributedQueryRunner createQueryRunner(Map<String, String> extra
  }
  }
 
+ public static BigQuery createBigQueryClient()
+ {
+ try {
+ InputStream jsonKey = new ByteArrayInputStream(Base64.getDecoder().decode(System.getProperty("bigquery.credentials-key")));
+ return BigQueryOptions.newBuilder()
+ .setCredentials(ServiceAccountCredentials.fromStream(jsonKey))
+ .build()
+ .getService();
+ }
+ catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
  public static Session createSession()
  {
  return testSessionBuilder()

diff --git a/...bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/...bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java
@@ -13,15 +13,22 @@
  */
 package io.prestosql.plugin.bigquery;
 
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.JobId;
+import com.google.cloud.bigquery.JobInfo;
+import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.common.collect.ImmutableMap;
 import io.prestosql.testing.AbstractTestIntegrationSmokeTest;
 import io.prestosql.testing.MaterializedResult;
 import io.prestosql.testing.QueryRunner;
 import org.testng.annotations.Test;
 
+import static io.prestosql.plugin.bigquery.BigQueryQueryRunner.createBigQueryClient;
 import static io.prestosql.spi.type.VarcharType.VARCHAR;
 import static io.prestosql.testing.MaterializedResult.resultBuilder;
 import static io.prestosql.testing.assertions.Assert.assertEquals;
+import static java.lang.String.format;
 
 @Test
 public class TestBigQueryIntegrationSmokeTest
@@ -51,4 +58,44 @@ public void testDescribeTable()
  MaterializedResult actualColumns = computeActual("DESCRIBE orders");
  assertEquals(actualColumns, expectedColumns);
  }
+
+ @Test(enabled = false)
+ public void testSelectFromHourlyPartitionedTable()
+ {
+ BigQuery client = createBigQueryClient();
+
+ executeBigQuerySql(client, "DROP TABLE IF EXISTS test.hourly_partitioned");
+ executeBigQuerySql(client, "CREATE TABLE test.hourly_partitioned (value INT64, ts TIMESTAMP) PARTITION BY TIMESTAMP_TRUNC(ts, HOUR)");
+ executeBigQuerySql(client, "INSERT INTO test.hourly_partitioned (value, ts) VALUES (1000, '2018-01-01 10:00:00')");
+
+ MaterializedResult actualValues = computeActual("SELECT COUNT(1) FROM test.hourly_partitioned");
+
+ assertEquals((long) actualValues.getOnlyValue(), 1L);
+ }
+
+ private static void executeBigQuerySql(BigQuery bigquery, String query)
+ {
+ QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query)
+ .setUseLegacySql(false)
+ .build();
+
+ JobId jobId = JobId.of();
+ Job queryJob = bigquery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());
+
+ try {
+ queryJob = queryJob.waitFor();
+
+ if (queryJob == null) {
+ throw new RuntimeException(format("Job with uuid %s does not longer exists", jobId.getJob()));
+ }
+
+ if (queryJob.getStatus().getError() != null) {
+ throw new RuntimeException(format("Query '%s' failed: %s", query, queryJob.getStatus().getError()));
+ }
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ }
 }