/*
 * 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.io.gcp.bigquery;

import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createJobIdToken;
import static org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createTempTableReference;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.hasItem;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;

import com.google.api.services.bigquery.model.JobStatistics;
import com.google.api.services.bigquery.model.JobStatistics2;
import com.google.api.services.bigquery.model.Streamingbuffer;
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableReference;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import com.google.bigtable.v2.Mutation;
import com.google.protobuf.ByteString;
import java.io.IOException;
import java.io.Serializable;
import java.math.BigInteger;
import java.util.List;
import java.util.Set;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.extensions.protobuf.ByteStringCoder;
import org.apache.beam.sdk.extensions.protobuf.ProtoCoder;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.QueryPriority;
import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
import org.apache.beam.sdk.io.gcp.testing.FakeJobService;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.transforms.Select;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.SourceTestUtils;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayDataEvaluator;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.rules.TemporaryFolder;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.junit.runners.model.Statement;

/** Tests for {@link BigQueryIO#read}. */
@RunWith(JUnit4.class)
public class BigQueryIOReadTest implements Serializable {
  private transient PipelineOptions options;
  private transient TemporaryFolder testFolder = new TemporaryFolder();
  private transient TestPipeline p;

  @Rule
  public final transient TestRule folderThenPipeline =
      new TestRule() {
        @Override
        public Statement apply(final Statement base, final Description description) {
          // We need to set up the temporary folder, and then set up the TestPipeline based on the
          // chosen folder. Unfortunately, since rule evaluation order is unspecified and unrelated
          // to field order, and is separate from construction, that requires manually creating this
          // TestRule.
          Statement withPipeline =
              new Statement() {
                @Override
                public void evaluate() throws Throwable {
                  options = TestPipeline.testingPipelineOptions();
                  options.as(BigQueryOptions.class).setProject("project-id");
                  options
                      .as(BigQueryOptions.class)
                      .setTempLocation(testFolder.getRoot().getAbsolutePath());
                  p = TestPipeline.fromOptions(options);
                  p.apply(base, description).evaluate();
                }
              };
          return testFolder.apply(withPipeline, description);
        }
      };

  @Rule public transient ExpectedException thrown = ExpectedException.none();

  private FakeDatasetService fakeDatasetService = new FakeDatasetService();
  private FakeJobService fakeJobService = new FakeJobService();
  private FakeBigQueryServices fakeBqServices =
      new FakeBigQueryServices()
          .withDatasetService(fakeDatasetService)
          .withJobService(fakeJobService);

  private void checkReadTableObject(
      BigQueryIO.Read read, String project, String dataset, String table) {
    checkReadTableObjectWithValidate(read, project, dataset, table, true);
  }

  private void checkReadQueryObject(BigQueryIO.Read read, String query) {
    checkReadQueryObjectWithValidate(read, query, true);
  }

  private void checkTypedReadQueryObject(BigQueryIO.TypedRead read, String query, String kmsKey) {
    checkTypedReadQueryObjectWithValidate(read, query, kmsKey, true);
  }

  private void checkReadTableObjectWithValidate(
      BigQueryIO.Read read, String project, String dataset, String table, boolean validate) {
    assertEquals(project, read.getTable().getProjectId());
    assertEquals(dataset, read.getTable().getDatasetId());
    assertEquals(table, read.getTable().getTableId());
    assertNull(read.getQuery());
    assertEquals(validate, read.getValidate());
  }

  private void checkReadQueryObjectWithValidate(
      BigQueryIO.Read read, String query, boolean validate) {
    assertNull(read.getTable());
    assertEquals(query, read.getQuery().get());
    assertEquals(validate, read.getValidate());
  }

  private void checkTypedReadQueryObjectWithValidate(
      BigQueryIO.TypedRead read, String query, String kmsKey, boolean validate) {
    assertNull(read.getTable());
    assertEquals(query, read.getQuery().get());
    assertEquals(kmsKey, read.getKmsKey());
    assertEquals(validate, read.getValidate());
  }

  @Before
  public void setUp() throws IOException, InterruptedException {
    FakeDatasetService.setUp();
    BigQueryIO.clearCreatedTables();
    fakeDatasetService.createDataset("project-id", "dataset-id", "", "", null);
  }

  @Test
  public void testBuildTableBasedSource() {
    BigQueryIO.Read read = BigQueryIO.read().from("foo.com:project:somedataset.sometable");
    checkReadTableObject(read, "foo.com:project", "somedataset", "sometable");
  }

  @Test
  public void testBuildQueryBasedSource() {
    BigQueryIO.Read read = BigQueryIO.read().fromQuery("foo_query");
    checkReadQueryObject(read, "foo_query");
  }

  @Test
  public void testBuildTableBasedSourceWithoutValidation() {
    // This test just checks that using withoutValidation will not trigger object
    // construction errors.
    BigQueryIO.Read read =
        BigQueryIO.read().from("foo.com:project:somedataset.sometable").withoutValidation();
    checkReadTableObjectWithValidate(read, "foo.com:project", "somedataset", "sometable", false);
  }

  @Test
  public void testBuildQueryBasedSourceWithoutValidation() {
    // This test just checks that using withoutValidation will not trigger object
    // construction errors.
    BigQueryIO.Read read = BigQueryIO.read().fromQuery("some_query").withoutValidation();
    checkReadQueryObjectWithValidate(read, "some_query", false);
  }

  @Test
  public void testBuildTableBasedSourceWithDefaultProject() {
    BigQueryIO.Read read = BigQueryIO.read().from("somedataset.sometable");
    checkReadTableObject(read, null, "somedataset", "sometable");
  }

  @Test
  public void testBuildSourceWithTableReference() {
    TableReference table =
        new TableReference()
            .setProjectId("foo.com:project")
            .setDatasetId("somedataset")
            .setTableId("sometable");
    BigQueryIO.Read read = BigQueryIO.read().from(table);
    checkReadTableObject(read, "foo.com:project", "somedataset", "sometable");
  }

  @Test
  public void testBuildQueryBasedTypedReadSource() {
    BigQueryIO.TypedRead read =
        BigQueryIO.readTableRows().fromQuery("foo_query").withKmsKey("kms_key");
    checkTypedReadQueryObject(read, "foo_query", "kms_key");
  }

  @Test
  public void testValidateReadSetsDefaultProject() throws Exception {
    String tableId = "sometable";
    TableReference tableReference =
        new TableReference()
            .setProjectId("project-id")
            .setDatasetId("dataset-id")
            .setTableId(tableId);
    fakeDatasetService.createTable(
        new Table()
            .setTableReference(tableReference)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER")))));

    FakeBigQueryServices fakeBqServices =
        new FakeBigQueryServices()
            .withJobService(new FakeJobService())
            .withDatasetService(fakeDatasetService);

    List<TableRow> expected =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));
    fakeDatasetService.insertAll(tableReference, expected, null);

    TableReference tableRef = new TableReference().setDatasetId("dataset-id").setTableId(tableId);

    PCollection<KV<String, Long>> output =
        p.apply(BigQueryIO.read().from(tableRef).withTestServices(fakeBqServices))
            .apply(
                ParDo.of(
                    new DoFn<TableRow, KV<String, Long>>() {
                      @ProcessElement
                      public void processElement(ProcessContext c) throws Exception {
                        c.output(
                            KV.of(
                                (String) c.element().get("name"),
                                Long.valueOf((String) c.element().get("number"))));
                      }
                    }));
    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(
                KV.of("a", 1L),
                KV.of("b", 2L),
                KV.of("c", 3L),
                KV.of("d", 4L),
                KV.of("e", 5L),
                KV.of("f", 6L)));
    p.run();
  }

  @Test
  public void testBuildSourceWithTableAndFlatten() {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(
        "Invalid BigQueryIO.Read: Specifies a table with a result flattening preference,"
            + " which only applies to queries");
    p.apply(
        "ReadMyTable",
        BigQueryIO.read().from("foo.com:project:somedataset.sometable").withoutResultFlattening());
    p.run();
  }

  @Test
  public void testBuildSourceWithTableAndFlattenWithoutValidation() {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(
        "Invalid BigQueryIO.Read: Specifies a table with a result flattening preference,"
            + " which only applies to queries");
    p.apply(
        BigQueryIO.read()
            .from("foo.com:project:somedataset.sometable")
            .withoutValidation()
            .withoutResultFlattening());
    p.run();
  }

  @Test
  public void testBuildSourceWithTableAndSqlDialect() {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(
        "Invalid BigQueryIO.Read: Specifies a table with a SQL dialect preference,"
            + " which only applies to queries");
    p.apply(BigQueryIO.read().from("foo.com:project:somedataset.sometable").usingStandardSql());
    p.run();
  }

  @Test
  public void testReadFromTableWithoutTemplateCompatibility()
      throws IOException, InterruptedException {
    testReadFromTable(false, false);
  }

  @Test
  public void testReadFromTableWithTemplateCompatibility()
      throws IOException, InterruptedException {
    testReadFromTable(true, false);
  }

  @Test
  public void testReadTableRowsFromTableWithoutTemplateCompatibility()
      throws IOException, InterruptedException {
    testReadFromTable(false, true);
  }

  @Test
  public void testReadTableRowsFromTableWithTemplateCompatibility()
      throws IOException, InterruptedException {
    testReadFromTable(true, true);
  }

  private void testReadFromTable(boolean useTemplateCompatibility, boolean useReadTableRows)
      throws IOException, InterruptedException {
    Table sometable = new Table();
    sometable.setSchema(
        new TableSchema()
            .setFields(
                ImmutableList.of(
                    new TableFieldSchema().setName("name").setType("STRING"),
                    new TableFieldSchema().setName("number").setType("INTEGER"))));
    sometable.setTableReference(
        new TableReference()
            .setProjectId("non-executing-project")
            .setDatasetId("somedataset")
            .setTableId("sometable"));
    sometable.setNumBytes(1024L * 1024L);
    FakeDatasetService fakeDatasetService = new FakeDatasetService();
    fakeDatasetService.createDataset("non-executing-project", "somedataset", "", "", null);
    fakeDatasetService.createTable(sometable);

    List<TableRow> records =
        Lists.newArrayList(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L));
    fakeDatasetService.insertAll(sometable.getTableReference(), records, null);

    FakeBigQueryServices fakeBqServices =
        new FakeBigQueryServices()
            .withJobService(new FakeJobService())
            .withDatasetService(fakeDatasetService);

    PTransform<PBegin, PCollection<TableRow>> readTransform;
    if (useReadTableRows) {
      BigQueryIO.Read read =
          BigQueryIO.read()
              .from("non-executing-project:somedataset.sometable")
              .withTestServices(fakeBqServices)
              .withoutValidation();
      readTransform = useTemplateCompatibility ? read.withTemplateCompatibility() : read;
    } else {
      BigQueryIO.TypedRead<TableRow> read =
          BigQueryIO.readTableRows()
              .from("non-executing-project:somedataset.sometable")
              .withTestServices(fakeBqServices)
              .withoutValidation();
      readTransform = useTemplateCompatibility ? read.withTemplateCompatibility() : read;
    }
    PCollection<KV<String, Long>> output =
        p.apply(readTransform)
            .apply(
                ParDo.of(
                    new DoFn<TableRow, KV<String, Long>>() {
                      @ProcessElement
                      public void processElement(ProcessContext c) throws Exception {
                        c.output(
                            KV.of(
                                (String) c.element().get("name"),
                                Long.valueOf((String) c.element().get("number"))));
                      }
                    }));

    PAssert.that(output)
        .containsInAnyOrder(ImmutableList.of(KV.of("a", 1L), KV.of("b", 2L), KV.of("c", 3L)));
    p.run();
  }

  @Test
  public void testReadTableWithSchema() throws IOException, InterruptedException {
    // setup
    Table someTable = new Table();
    someTable.setSchema(
        new TableSchema()
            .setFields(
                ImmutableList.of(
                    new TableFieldSchema().setName("name").setType("STRING"),
                    new TableFieldSchema().setName("number").setType("INTEGER"))));
    someTable.setTableReference(
        new TableReference()
            .setProjectId("non-executing-project")
            .setDatasetId("schema_dataset")
            .setTableId("schema_table"));
    someTable.setNumBytes(1024L * 1024L);
    FakeDatasetService fakeDatasetService = new FakeDatasetService();
    fakeDatasetService.createDataset("non-executing-project", "schema_dataset", "", "", null);
    fakeDatasetService.createTable(someTable);

    List<TableRow> records =
        Lists.newArrayList(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L));

    fakeDatasetService.insertAll(someTable.getTableReference(), records, null);

    FakeBigQueryServices fakeBqServices =
        new FakeBigQueryServices()
            .withJobService(new FakeJobService())
            .withDatasetService(fakeDatasetService);

    // test
    BigQueryIO.TypedRead<TableRow> read =
        BigQueryIO.readTableRowsWithSchema()
            .from("non-executing-project:schema_dataset.schema_table")
            .withTestServices(fakeBqServices)
            .withoutValidation();

    PCollection<TableRow> bqRows = p.apply(read);

    Schema expectedSchema =
        Schema.of(
            Schema.Field.of("name", Schema.FieldType.STRING).withNullable(true),
            Schema.Field.of("number", Schema.FieldType.INT64).withNullable(true));
    assertEquals(expectedSchema, bqRows.getSchema());

    PCollection<Row> output = bqRows.apply(Select.fieldNames("name", "number"));
    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(
                Row.withSchema(expectedSchema).addValues("a", 1L).build(),
                Row.withSchema(expectedSchema).addValues("b", 2L).build(),
                Row.withSchema(expectedSchema).addValues("c", 3L).build()));

    p.run();
  }

  @Test
  public void testBuildSourceDisplayDataTable() {
    String tableSpec = "project:dataset.tableid";

    BigQueryIO.Read read =
        BigQueryIO.read()
            .from(tableSpec)
            .withoutResultFlattening()
            .usingStandardSql()
            .withoutValidation();

    DisplayData displayData = DisplayData.from(read);

    assertThat(displayData, hasDisplayItem("table", tableSpec));
    assertThat(displayData, hasDisplayItem("flattenResults", false));
    assertThat(displayData, hasDisplayItem("useLegacySql", false));
    assertThat(displayData, hasDisplayItem("validation", false));
  }

  @Test
  public void testBuildSourceDisplayDataQuery() {
    BigQueryIO.Read read =
        BigQueryIO.read()
            .fromQuery("myQuery")
            .withoutResultFlattening()
            .usingStandardSql()
            .withoutValidation();

    DisplayData displayData = DisplayData.from(read);

    assertThat(displayData, hasDisplayItem("query", "myQuery"));
    assertThat(displayData, hasDisplayItem("flattenResults", false));
    assertThat(displayData, hasDisplayItem("useLegacySql", false));
    assertThat(displayData, hasDisplayItem("validation", false));
  }

  @Test
  public void testTableSourcePrimitiveDisplayData() throws IOException, InterruptedException {
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    BigQueryIO.Read read =
        BigQueryIO.read()
            .from("project:dataset.tableId")
            .withTestServices(
                new FakeBigQueryServices()
                    .withDatasetService(new FakeDatasetService())
                    .withJobService(new FakeJobService()))
            .withoutValidation();

    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
    assertThat(
        "BigQueryIO.Read should include the table spec in its primitive display data",
        displayData,
        hasItem(hasDisplayItem("table")));
  }

  @Test
  public void testQuerySourcePrimitiveDisplayData() throws IOException, InterruptedException {
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    BigQueryIO.Read read =
        BigQueryIO.read()
            .fromQuery("foobar")
            .withTestServices(
                new FakeBigQueryServices()
                    .withDatasetService(new FakeDatasetService())
                    .withJobService(new FakeJobService()))
            .withoutValidation();

    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
    assertThat(
        "BigQueryIO.Read should include the query in its primitive display data",
        displayData,
        hasItem(hasDisplayItem("query")));
  }

  @Test
  public void testBigQueryIOGetName() {
    assertEquals("BigQueryIO.Read", BigQueryIO.read().from("somedataset.sometable").getName());
  }

  @Test
  public void testBigQueryTableSourceInitSplit() throws Exception {
    List<TableRow> expected =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));

    TableReference table = BigQueryHelpers.parseTableSpec("project:data_set.table_name");
    fakeDatasetService.createDataset("project", "data_set", "", "", null);
    fakeDatasetService.createTable(
        new Table()
            .setTableReference(table)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER")))));
    fakeDatasetService.insertAll(table, expected, null);

    String stepUuid = "testStepUuid";
    BoundedSource<TableRow> bqSource =
        BigQueryTableSourceDef.create(fakeBqServices, ValueProvider.StaticValueProvider.of(table))
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    PipelineOptions options = PipelineOptionsFactory.create();
    options.setTempLocation(testFolder.getRoot().getAbsolutePath());
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");

    List<TableRow> read =
        convertStringsToLong(
            SourceTestUtils.readFromSplitsOfSource(bqSource, 0L /* ignored */, options));
    assertThat(read, containsInAnyOrder(Iterables.toArray(expected, TableRow.class)));

    List<? extends BoundedSource<TableRow>> sources = bqSource.split(100, options);
    assertEquals(2, sources.size());
    // Simulate a repeated call to split(), like a Dataflow worker will sometimes do.
    sources = bqSource.split(200, options);
    assertEquals(2, sources.size());

    // A repeated call to split() should not have caused a duplicate extract job.
    assertEquals(1, fakeJobService.getNumExtractJobCalls());
  }

  @Test
  public void testEstimatedSizeWithoutStreamingBuffer() throws Exception {
    List<TableRow> data =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));

    TableReference table = BigQueryHelpers.parseTableSpec("project:data_set.table_name");
    fakeDatasetService.createDataset("project", "data_set", "", "", null);
    fakeDatasetService.createTable(
        new Table()
            .setTableReference(table)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER")))));
    fakeDatasetService.insertAll(table, data, null);

    String stepUuid = "testStepUuid";
    BoundedSource<TableRow> bqSource =
        BigQueryTableSourceDef.create(fakeBqServices, ValueProvider.StaticValueProvider.of(table))
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    PipelineOptions options = PipelineOptionsFactory.create();
    assertEquals(108, bqSource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testEstimatedSizeWithStreamingBuffer() throws Exception {
    List<TableRow> data =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));

    TableReference table = BigQueryHelpers.parseTableSpec("project:data_set.table_name");
    fakeDatasetService.createDataset("project", "data_set", "", "", null);
    fakeDatasetService.createTable(
        new Table()
            .setTableReference(table)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER"))))
            .setStreamingBuffer(new Streamingbuffer().setEstimatedBytes(BigInteger.valueOf(10))));
    fakeDatasetService.insertAll(table, data, null);

    String stepUuid = "testStepUuid";
    BoundedSource<TableRow> bqSource =
        BigQueryTableSourceDef.create(fakeBqServices, ValueProvider.StaticValueProvider.of(table))
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    PipelineOptions options = PipelineOptionsFactory.create();
    assertEquals(118, bqSource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testBigQueryQuerySourceEstimatedSize() throws Exception {

    String queryString = "fake query string";

    PipelineOptions options = PipelineOptionsFactory.create();
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");
    String stepUuid = "testStepUuid";

    BigQuerySourceBase<TableRow> bqSource =
        BigQueryQuerySourceDef.create(
                fakeBqServices,
                ValueProvider.StaticValueProvider.of(queryString),
                true, /* flattenResults */
                true, /* useLegacySql */
                QueryPriority.BATCH,
                null,
                null)
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    fakeJobService.expectDryRunQuery(
        bqOptions.getProject(),
        queryString,
        new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L)));

    assertEquals(100, bqSource.getEstimatedSizeBytes(bqOptions));
  }

  @Test
  public void testBigQueryQuerySourceInitSplit() throws Exception {

    PipelineOptions options = PipelineOptionsFactory.create();
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");

    TableReference sourceTableRef = BigQueryHelpers.parseTableSpec("project:dataset.table");

    fakeDatasetService.createDataset(
        sourceTableRef.getProjectId(),
        sourceTableRef.getDatasetId(),
        "asia-northeast1",
        "Fake plastic tree^H^H^H^Htables",
        null);

    fakeDatasetService.createTable(
        new Table().setTableReference(sourceTableRef).setLocation("asia-northeast1"));

    Table queryResultTable =
        new Table()
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER"))));

    List<TableRow> expected =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable, expected);

    String stepUuid = "testStepUuid";

    TableReference tempTableReference =
        createTempTableReference(
            bqOptions.getProject(), createJobIdToken(options.getJobName(), stepUuid));

    fakeJobService.expectDryRunQuery(
        bqOptions.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(100L)
                    .setReferencedTables(ImmutableList.of(sourceTableRef, tempTableReference))));

    BoundedSource<TableRow> bqSource =
        BigQueryQuerySourceDef.create(
                fakeBqServices,
                ValueProvider.StaticValueProvider.of(encodedQuery),
                true /* flattenResults */,
                true /* useLegacySql */,
                QueryPriority.BATCH,
                null,
                null)
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    options.setTempLocation(testFolder.getRoot().getAbsolutePath());

    List<TableRow> read =
        convertStringsToLong(
            SourceTestUtils.readFromSplitsOfSource(bqSource, 0L /* ignored */, options));
    assertThat(read, containsInAnyOrder(Iterables.toArray(expected, TableRow.class)));

    List<? extends BoundedSource<TableRow>> sources = bqSource.split(100, options);
    assertEquals(2, sources.size());
  }

  /**
   * This test simulates the scenario where the SQL text which is executed by the query job doesn't
   * by itself refer to any tables (e.g. "SELECT 17 AS value"), and thus there are no referenced
   * tables when the dry run of the query is performed.
   */
  @Test
  public void testBigQueryQuerySourceInitSplit_NoReferencedTables() throws Exception {

    PipelineOptions options = PipelineOptionsFactory.create();
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");

    Table queryResultTable =
        new Table()
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER"))));

    List<TableRow> expected =
        ImmutableList.of(
            new TableRow().set("name", "a").set("number", 1L),
            new TableRow().set("name", "b").set("number", 2L),
            new TableRow().set("name", "c").set("number", 3L),
            new TableRow().set("name", "d").set("number", 4L),
            new TableRow().set("name", "e").set("number", 5L),
            new TableRow().set("name", "f").set("number", 6L));

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable, expected);

    String stepUuid = "testStepUuid";

    fakeJobService.expectDryRunQuery(
        bqOptions.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(100L)
                    .setReferencedTables(ImmutableList.of())));

    BoundedSource<TableRow> bqSource =
        BigQueryQuerySourceDef.create(
                fakeBqServices,
                ValueProvider.StaticValueProvider.of(encodedQuery),
                true /* flattenResults */,
                true /* useLegacySql */,
                QueryPriority.BATCH,
                null,
                null)
            .toSource(stepUuid, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE);

    options.setTempLocation(testFolder.getRoot().getAbsolutePath());

    List<TableRow> read =
        convertStringsToLong(
            SourceTestUtils.readFromSplitsOfSource(bqSource, 0L /* ignored */, options));
    assertThat(read, containsInAnyOrder(Iterables.toArray(expected, TableRow.class)));

    List<? extends BoundedSource<TableRow>> sources = bqSource.split(100, options);
    assertEquals(2, sources.size());
  }

  @Test
  public void testPassThroughThenCleanup() throws Exception {

    PCollection<Integer> output =
        p.apply(Create.of(1, 2, 3))
            .apply(
                new PassThroughThenCleanup<>(
                    new PassThroughThenCleanup.CleanupOperation() {
                      @Override
                      void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception {
                        // no-op
                      }
                    },
                    p.apply("Create1", Create.of("")).apply(View.asSingleton())));

    PAssert.that(output).containsInAnyOrder(1, 2, 3);

    p.run();
  }

  @Test
  public void testPassThroughThenCleanupExecuted() throws Exception {

    p.apply(Create.empty(VarIntCoder.of()))
        .apply(
            new PassThroughThenCleanup<>(
                new PassThroughThenCleanup.CleanupOperation() {
                  @Override
                  void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception {
                    throw new RuntimeException("cleanup executed");
                  }
                },
                p.apply("Create1", Create.of("")).apply(View.asSingleton())));

    thrown.expect(RuntimeException.class);
    thrown.expectMessage("cleanup executed");

    p.run();
  }

  @Test
  public void testRuntimeOptionsNotCalledInApplyInputTable() {
    BigQueryIO.Read read = BigQueryIO.read().from(p.newProvider("")).withoutValidation();
    // Test that this doesn't throw.
    DisplayData.from(read);
  }

  @Test
  public void testRuntimeOptionsNotCalledInApplyInputQuery() {
    BigQueryIO.Read read = BigQueryIO.read().fromQuery(p.newProvider("")).withoutValidation();
    // Test that this doesn't throw.
    DisplayData.from(read);
  }

  List<TableRow> convertStringsToLong(List<TableRow> toConvert) {
    // The numbers come back as String after JSON serialization. Change them back to
    // longs so that we can assert the output.
    List<TableRow> converted = Lists.newArrayList();
    for (TableRow entry : toConvert) {
      TableRow convertedEntry = entry.clone();
      convertedEntry.set("number", Long.parseLong((String) convertedEntry.get("number")));
      converted.add(convertedEntry);
    }
    return converted;
  }

  @Test
  public void testCoderInference() {
    // Lambdas erase too much type information - use an anonymous class here.
    SerializableFunction<SchemaAndRecord, KV<ByteString, Mutation>> parseFn =
        new SerializableFunction<SchemaAndRecord, KV<ByteString, Mutation>>() {
          @Override
          public KV<ByteString, Mutation> apply(SchemaAndRecord input) {
            return null;
          }
        };

    assertEquals(
        KvCoder.of(ByteStringCoder.of(), ProtoCoder.of(Mutation.class)),
        BigQueryIO.read(parseFn).inferCoder(CoderRegistry.createDefault()));
  }
}