| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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 org.apache.beam.sdk.extensions.sql.meta.provider.bigquery; |
| |
| import static org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils.parseTimestampWithUTCTimeZone; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.BOOLEAN; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.BYTE; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.DOUBLE; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.FLOAT; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.INT16; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.INT32; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.INT64; |
| import static org.apache.beam.sdk.schemas.Schema.FieldType.STRING; |
| import static org.hamcrest.Matchers.containsInAnyOrder; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertThat; |
| |
| import java.io.Serializable; |
| import java.util.Arrays; |
| import java.util.List; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.PipelineResult; |
| import org.apache.beam.sdk.PipelineResult.State; |
| import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv; |
| import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; |
| import org.apache.beam.sdk.extensions.sql.impl.schema.BeamPCollectionTable; |
| import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; |
| import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider; |
| import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; |
| import org.apache.beam.sdk.io.gcp.bigquery.TestBigQuery; |
| import org.apache.beam.sdk.schemas.Schema; |
| import org.apache.beam.sdk.schemas.Schema.FieldType; |
| import org.apache.beam.sdk.testing.PAssert; |
| import org.apache.beam.sdk.testing.TestPipeline; |
| import org.apache.beam.sdk.transforms.Create; |
| import org.apache.beam.sdk.transforms.SerializableFunctions; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.Row; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap; |
| import org.joda.time.Duration; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** Integration tests form writing to BigQuery with Beam SQL. */ |
| @RunWith(JUnit4.class) |
| public class BigQueryReadWriteIT implements Serializable { |
| private static final Schema SOURCE_SCHEMA = |
| Schema.builder() |
| .addNullableField("id", INT64) |
| .addNullableField("name", STRING) |
| .addNullableField("arr", FieldType.array(STRING)) |
| .build(); |
| |
| private static final Schema SOURCE_SCHEMA_TWO = |
| Schema.builder() |
| .addNullableField("c_bigint", INT64) |
| .addNullableField("c_tinyint", BYTE) |
| .addNullableField("c_smallint", INT16) |
| .addNullableField("c_integer", INT32) |
| .addNullableField("c_float", FLOAT) |
| .addNullableField("c_double", DOUBLE) |
| .addNullableField("c_boolean", BOOLEAN) |
| .addNullableField("c_timestamp", CalciteUtils.TIMESTAMP) |
| .addNullableField("c_varchar", STRING) |
| .addNullableField("c_char", STRING) |
| .addNullableField("c_arr", FieldType.array(STRING)) |
| .build(); |
| |
| @Rule public transient TestPipeline pipeline = TestPipeline.create(); |
| @Rule public transient TestPipeline readPipeline = TestPipeline.create(); |
| @Rule public transient TestBigQuery bigQuery = TestBigQuery.create(SOURCE_SCHEMA); |
| @Rule public transient TestBigQuery bigQueryTestingTypes = TestBigQuery.create(SOURCE_SCHEMA_TWO); |
| |
| @Test |
| public void testSQLRead() { |
| BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new BigQueryTableProvider()); |
| |
| String createTableStatement = |
| "CREATE EXTERNAL TABLE TEST( \n" |
| + " c_bigint BIGINT, \n" |
| + " c_tinyint TINYINT, \n" |
| + " c_smallint SMALLINT, \n" |
| + " c_integer INTEGER, \n" |
| + " c_float FLOAT, \n" |
| + " c_double DOUBLE, \n" |
| + " c_boolean BOOLEAN, \n" |
| + " c_timestamp TIMESTAMP, \n" |
| + " c_varchar VARCHAR, \n " |
| + " c_char CHAR, \n" |
| + " c_arr ARRAY<VARCHAR> \n" |
| + ") \n" |
| + "TYPE 'bigquery' \n" |
| + "LOCATION '" |
| + bigQueryTestingTypes.tableSpec() |
| + "'"; |
| sqlEnv.executeDdl(createTableStatement); |
| |
| String insertStatement = |
| "INSERT INTO TEST VALUES (" |
| + "9223372036854775807, " |
| + "127, " |
| + "32767, " |
| + "2147483647, " |
| + "1.0, " |
| + "1.0, " |
| + "TRUE, " |
| + "TIMESTAMP '2018-05-28 20:17:40.123', " |
| + "'varchar', " |
| + "'char', " |
| + "ARRAY['123', '456']" |
| + ")"; |
| |
| sqlEnv.parseQuery(insertStatement); |
| BeamSqlRelUtils.toPCollection(pipeline, sqlEnv.parseQuery(insertStatement)); |
| pipeline.run().waitUntilFinish(Duration.standardMinutes(5)); |
| |
| String selectTableStatement = "SELECT * FROM TEST"; |
| PCollection<Row> output = |
| BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement)); |
| |
| PAssert.that(output) |
| .containsInAnyOrder( |
| row( |
| SOURCE_SCHEMA_TWO, |
| 9223372036854775807L, |
| (byte) 127, |
| (short) 32767, |
| 2147483647, |
| (float) 1.0, |
| 1.0, |
| true, |
| parseTimestampWithUTCTimeZone("2018-05-28 20:17:40.123"), |
| "varchar", |
| "char", |
| Arrays.asList("123", "456"))); |
| PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5)); |
| assertEquals(state, State.DONE); |
| } |
| |
| @Test |
| public void testSQLTypes() { |
| BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new BigQueryTableProvider()); |
| |
| String createTableStatement = |
| "CREATE EXTERNAL TABLE TEST( \n" |
| + " c_bigint BIGINT, \n" |
| + " c_tinyint TINYINT, \n" |
| + " c_smallint SMALLINT, \n" |
| + " c_integer INTEGER, \n" |
| + " c_float FLOAT, \n" |
| + " c_double DOUBLE, \n" |
| + " c_boolean BOOLEAN, \n" |
| + " c_timestamp TIMESTAMP, \n" |
| + " c_varchar VARCHAR, \n " |
| + " c_char CHAR, \n" |
| + " c_arr ARRAY<VARCHAR> \n" |
| + ") \n" |
| + "TYPE 'bigquery' \n" |
| + "LOCATION '" |
| + bigQueryTestingTypes.tableSpec() |
| + "'"; |
| sqlEnv.executeDdl(createTableStatement); |
| |
| String insertStatement = |
| "INSERT INTO TEST VALUES (" |
| + "9223372036854775807, " |
| + "127, " |
| + "32767, " |
| + "2147483647, " |
| + "1.0, " |
| + "1.0, " |
| + "TRUE, " |
| + "TIMESTAMP '2018-05-28 20:17:40.123', " |
| + "'varchar', " |
| + "'char', " |
| + "ARRAY['123', '456']" |
| + ")"; |
| sqlEnv.parseQuery(insertStatement); |
| BeamSqlRelUtils.toPCollection(pipeline, sqlEnv.parseQuery(insertStatement)); |
| pipeline.run().waitUntilFinish(Duration.standardMinutes(5)); |
| |
| assertThat( |
| bigQueryTestingTypes.getFlatJsonRows(SOURCE_SCHEMA_TWO), |
| containsInAnyOrder( |
| row( |
| SOURCE_SCHEMA_TWO, |
| 9223372036854775807L, |
| (byte) 127, |
| (short) 32767, |
| 2147483647, |
| (float) 1.0, |
| 1.0, |
| true, |
| parseTimestampWithUTCTimeZone("2018-05-28 20:17:40.123"), |
| "varchar", |
| "char", |
| Arrays.asList("123", "456")))); |
| } |
| |
| @Test |
| public void testInsertSelect() throws Exception { |
| BeamSqlEnv sqlEnv = |
| BeamSqlEnv.inMemory( |
| readOnlyTableProvider( |
| pipeline, |
| "ORDERS_IN_MEMORY", |
| row(SOURCE_SCHEMA, 1L, "foo", Arrays.asList("111", "aaa")), |
| row(SOURCE_SCHEMA, 2L, "bar", Arrays.asList("222", "bbb")), |
| row(SOURCE_SCHEMA, 3L, "baz", Arrays.asList("333", "ccc"))), |
| new BigQueryTableProvider()); |
| |
| String createTableStatement = |
| "CREATE EXTERNAL TABLE ORDERS_BQ( \n" |
| + " id BIGINT, \n" |
| + " name VARCHAR, \n " |
| + " arr ARRAY<VARCHAR> \n" |
| + ") \n" |
| + "TYPE 'bigquery' \n" |
| + "LOCATION '" |
| + bigQuery.tableSpec() |
| + "'"; |
| sqlEnv.executeDdl(createTableStatement); |
| |
| String insertStatement = |
| "INSERT INTO ORDERS_BQ \n" |
| + " SELECT \n" |
| + " id as `id`, \n" |
| + " name as `name`, \n" |
| + " arr as `arr` \n" |
| + " FROM ORDERS_IN_MEMORY"; |
| BeamSqlRelUtils.toPCollection(pipeline, sqlEnv.parseQuery(insertStatement)); |
| |
| pipeline.run().waitUntilFinish(Duration.standardMinutes(5)); |
| |
| List<Row> allJsonRows = bigQuery.getFlatJsonRows(SOURCE_SCHEMA); |
| |
| assertThat( |
| allJsonRows, |
| containsInAnyOrder( |
| row(SOURCE_SCHEMA, 1L, "foo", Arrays.asList("111", "aaa")), |
| row(SOURCE_SCHEMA, 2L, "bar", Arrays.asList("222", "bbb")), |
| row(SOURCE_SCHEMA, 3L, "baz", Arrays.asList("333", "ccc")))); |
| } |
| |
| private TableProvider readOnlyTableProvider(Pipeline pipeline, String tableName, Row... rows) { |
| |
| return new ReadOnlyTableProvider( |
| "PCOLLECTION", |
| ImmutableMap.of(tableName, new BeamPCollectionTable(createPCollection(pipeline, rows)))); |
| } |
| |
| private PCollection<Row> createPCollection(Pipeline pipeline, Row... rows) { |
| return pipeline.apply( |
| Create.of(Arrays.asList(rows)) |
| .withSchema( |
| SOURCE_SCHEMA, SerializableFunctions.identity(), SerializableFunctions.identity())); |
| } |
| |
| private Row row(Schema schema, Object... values) { |
| return Row.withSchema(schema).addValues(values).build(); |
| } |
| } |
