/*
 * 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.BigQueryResourceNaming.createTempTableReference;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.withSettings;

import com.google.api.services.bigquery.model.JobStatistics;
import com.google.api.services.bigquery.model.JobStatistics2;
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.cloud.bigquery.storage.v1.AvroRows;
import com.google.cloud.bigquery.storage.v1.AvroSchema;
import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
import com.google.cloud.bigquery.storage.v1.DataFormat;
import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
import com.google.cloud.bigquery.storage.v1.ReadSession;
import com.google.cloud.bigquery.storage.v1.ReadStream;
import com.google.cloud.bigquery.storage.v1.StreamStats;
import com.google.cloud.bigquery.storage.v1.StreamStats.Progress;
import com.google.protobuf.ByteString;
import java.io.ByteArrayOutputStream;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import org.apache.avro.Schema;
import org.apache.avro.SchemaBuilder;
import org.apache.avro.generic.GenericData.Record;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.Encoder;
import org.apache.avro.io.EncoderFactory;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.KvCoder;
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.TableRowParser;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.QueryPriority;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryResourceNaming.JobType;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient;
import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices.FakeBigQueryServerStream;
import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
import org.apache.beam.sdk.io.gcp.testing.FakeJobService;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
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.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#readTableRows()} using {@link Method#DIRECT_READ}. */
@RunWith(JUnit4.class)
public class BigQueryIOStorageQueryTest {

  private transient BigQueryOptions options;
  private transient TemporaryFolder testFolder = new TemporaryFolder();
  private transient TestPipeline p;

  @Rule
  public final transient TestRule folderThenPipeline =
      new TestRule() {
        @Override
        public Statement apply(Statement base, 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().as(BigQueryOptions.class);
                  options.setProject("project-id");
                  if (description.getAnnotations().stream()
                      .anyMatch(a -> a.annotationType().equals(ProjectOverride.class))) {
                    options.as(BigQueryOptions.class).setBigQueryProject("bigquery-project-id");
                  }
                  options.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 fakeBigQueryServices =
      new FakeBigQueryServices()
          .withDatasetService(fakeDatasetService)
          .withJobService(fakeJobService);

  @Before
  public void setUp() throws Exception {
    FakeDatasetService.setUp();
  }

  private static final String DEFAULT_QUERY = "SELECT * FROM `dataset.table` LIMIT 1";

  @Test
  public void testDefaultQueryBasedSource() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead();
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertTrue(typedRead.getValidate());
    assertTrue(typedRead.getFlattenResults());
    assertTrue(typedRead.getUseLegacySql());
    assertNull(typedRead.getQueryPriority());
    assertNull(typedRead.getQueryLocation());
    assertNull(typedRead.getKmsKey());
    assertFalse(typedRead.getWithTemplateCompatibility());
  }

  @Test
  public void testQueryBasedSourceWithCustomQuery() throws Exception {
    TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .fromQuery("SELECT * FROM `google.com:project.dataset.table`")
            .withCoder(TableRowJsonCoder.of());
    checkTypedReadQueryObject(typedRead, "SELECT * FROM `google.com:project.dataset.table`");
  }

  @Test
  public void testQueryBasedSourceWithoutValidation() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withoutValidation();
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertFalse(typedRead.getValidate());
  }

  @Test
  public void testQueryBasedSourceWithoutResultFlattening() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withoutResultFlattening();
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertFalse(typedRead.getFlattenResults());
  }

  @Test
  public void testQueryBasedSourceWithStandardSql() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().usingStandardSql();
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertFalse(typedRead.getUseLegacySql());
  }

  @Test
  public void testQueryBasedSourceWithPriority() throws Exception {
    TypedRead<TableRow> typedRead =
        getDefaultTypedRead().withQueryPriority(QueryPriority.INTERACTIVE);
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertEquals(QueryPriority.INTERACTIVE, typedRead.getQueryPriority());
  }

  @Test
  public void testQueryBasedSourceWithQueryLocation() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withQueryLocation("US");
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertEquals("US", typedRead.getQueryLocation());
  }

  @Test
  public void testQueryBasedSourceWithKmsKey() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withKmsKey("kms_key");
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertEquals("kms_key", typedRead.getKmsKey());
  }

  @Test
  public void testQueryBasedSourceWithTemplateCompatibility() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withTemplateCompatibility();
    checkTypedReadQueryObject(typedRead, DEFAULT_QUERY);
    assertTrue(typedRead.getWithTemplateCompatibility());
  }

  private TypedRead<TableRow> getDefaultTypedRead() {
    return BigQueryIO.read(new TableRowParser())
        .fromQuery(DEFAULT_QUERY)
        .withCoder(TableRowJsonCoder.of())
        .withMethod(Method.DIRECT_READ);
  }

  private void checkTypedReadQueryObject(TypedRead<?> typedRead, String query) {
    assertNull(typedRead.getTable());
    assertEquals(query, typedRead.getQuery().get());
  }

  @Test
  public void testBuildQueryBasedSourceWithSelectedFields() throws Exception {
    TypedRead<TableRow> typedRead =
        getDefaultTypedRead().withSelectedFields(Lists.newArrayList("a"));
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(
        "Invalid BigQueryIO.Read: Specifies selected fields, "
            + "which only applies when reading from a table");
    p.apply(typedRead);
    p.run();
  }

  @Test
  public void testBuildQueryBasedSourceWithRowRestriction() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead().withRowRestriction("a > 5");
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage(
        "Invalid BigQueryIO.Read: Specifies row restriction, "
            + "which only applies when reading from a table");
    p.apply(typedRead);
    p.run();
  }

  @Test
  public void testDisplayData() throws Exception {
    TypedRead<TableRow> typedRead = getDefaultTypedRead();
    DisplayData displayData = DisplayData.from(typedRead);
    assertThat(displayData, hasDisplayItem("query", DEFAULT_QUERY));
  }

  @Test
  public void testName() {
    assertEquals("BigQueryIO.TypedRead", getDefaultTypedRead().getName());
  }

  @Test
  public void testCoderInference() {
    SerializableFunction<SchemaAndRecord, KV<ByteString, ReadSession>> parseFn =
        new SerializableFunction<SchemaAndRecord, KV<ByteString, ReadSession>>() {
          @Override
          public KV<ByteString, ReadSession> apply(SchemaAndRecord input) {
            return null;
          }
        };

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

  @Test
  public void testQuerySourceEstimatedSize() throws Exception {

    String fakeQuery = "fake query text";

    fakeJobService.expectDryRunQuery(
        options.getProject(),
        fakeQuery,
        new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(125L)));

    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            /* stepUuid = */ "stepUuid",
            ValueProvider.StaticValueProvider.of(fakeQuery),
            /* flattenResults = */ true,
            /* useLegacySql = */ true,
            /* priority = */ QueryPriority.INTERACTIVE,
            /* location = */ null,
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            fakeBigQueryServices);

    assertEquals(125L, querySource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testQuerySourceInitialSplit() throws Exception {
    doQuerySourceInitialSplit(1024L, 1024, 50);
  }

  @Test
  public void testQuerySourceInitialSplit_MinSplitCount() throws Exception {
    doQuerySourceInitialSplit(1024L * 1024L, 10, 1);
  }

  @Test
  public void testQuerySourceInitialSplit_MaxSplitCount() throws Exception {
    doQuerySourceInitialSplit(10, 10_000, 200);
  }

  private void doQuerySourceInitialSplit(
      long bundleSize, int requestedStreamCount, int expectedStreamCount) throws Exception {

    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"))))
            .setNumBytes(1024L * 1024L);

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable);

    fakeJobService.expectDryRunQuery(
        options.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(1024L * 1024L)
                    .setReferencedTables(ImmutableList.of(sourceTableRef))));

    String stepUuid = "testStepUuid";

    TableReference tempTableReference =
        createTempTableReference(
            options.getProject(),
            BigQueryResourceNaming.createJobIdPrefix(options.getJobName(), stepUuid, JobType.QUERY),
            Optional.empty());

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/" + options.getProject())
            .setReadSession(
                ReadSession.newBuilder()
                    .setTable(BigQueryHelpers.toTableResourceName(tempTableReference)))
            .setMaxStreamCount(requestedStreamCount)
            .build();

    Schema sessionSchema =
        SchemaBuilder.record("__root__")
            .fields()
            .name("name")
            .type()
            .nullable()
            .stringType()
            .noDefault()
            .name("number")
            .type()
            .nullable()
            .longType()
            .noDefault()
            .endRecord();

    ReadSession.Builder builder =
        ReadSession.newBuilder()
            .setAvroSchema(AvroSchema.newBuilder().setSchema(sessionSchema.toString()))
            .setDataFormat(DataFormat.AVRO);
    for (int i = 0; i < expectedStreamCount; i++) {
      builder.addStreams(ReadStream.newBuilder().setName("stream-" + i));
    }

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.createReadSession(expectedRequest)).thenReturn(builder.build());

    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            stepUuid,
            ValueProvider.StaticValueProvider.of(encodedQuery),
            /* flattenResults = */ true,
            /* useLegacySql = */ true,
            /* priority = */ QueryPriority.BATCH,
            /* location = */ null,
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withJobService(fakeJobService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = querySource.split(bundleSize, options);
    assertEquals(expectedStreamCount, 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 testQuerySourceInitialSplit_NoReferencedTables() throws Exception {

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

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable);

    fakeJobService.expectDryRunQuery(
        options.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(1024L * 1024L)
                    .setReferencedTables(ImmutableList.of())));

    String stepUuid = "testStepUuid";

    TableReference tempTableReference =
        createTempTableReference(
            options.getProject(),
            BigQueryResourceNaming.createJobIdPrefix(options.getJobName(), stepUuid, JobType.QUERY),
            Optional.empty());

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/" + options.getProject())
            .setReadSession(
                ReadSession.newBuilder()
                    .setTable(BigQueryHelpers.toTableResourceName(tempTableReference)))
            .setMaxStreamCount(1024)
            .build();

    Schema sessionSchema =
        SchemaBuilder.record("__root__")
            .fields()
            .name("name")
            .type()
            .nullable()
            .stringType()
            .noDefault()
            .name("number")
            .type()
            .nullable()
            .longType()
            .noDefault()
            .endRecord();

    ReadSession.Builder builder =
        ReadSession.newBuilder()
            .setAvroSchema(AvroSchema.newBuilder().setSchema(sessionSchema.toString()))
            .setDataFormat(DataFormat.AVRO);
    for (int i = 0; i < 1024; i++) {
      builder.addStreams(ReadStream.newBuilder().setName("stream-" + i));
    }

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.createReadSession(expectedRequest)).thenReturn(builder.build());

    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            stepUuid,
            ValueProvider.StaticValueProvider.of(encodedQuery),
            /* flattenResults = */ true,
            /* useLegacySql = */ true,
            /* priority = */ QueryPriority.BATCH,
            /* location = */ null,
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withJobService(fakeJobService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = querySource.split(1024, options);
    assertEquals(1024, sources.size());
  }

  private static final String AVRO_SCHEMA_STRING =
      "{\"namespace\": \"example.avro\",\n"
          + " \"type\": \"record\",\n"
          + " \"name\": \"RowRecord\",\n"
          + " \"fields\": [\n"
          + "     {\"name\": \"name\", \"type\": [\"null\", \"string\"]},\n"
          + "     {\"name\": \"number\", \"type\": [\"null\", \"long\"]}\n"
          + " ]\n"
          + "}";

  private static final Schema AVRO_SCHEMA = new Schema.Parser().parse(AVRO_SCHEMA_STRING);

  private static final TableSchema TABLE_SCHEMA =
      new TableSchema()
          .setFields(
              ImmutableList.of(
                  new TableFieldSchema().setName("name").setType("STRING"),
                  new TableFieldSchema().setName("number").setType("INTEGER")));

  private static GenericRecord createRecord(String name, long number, Schema schema) {
    GenericRecord genericRecord = new Record(schema);
    genericRecord.put("name", name);
    genericRecord.put("number", number);
    return genericRecord;
  }

  private static final EncoderFactory ENCODER_FACTORY = EncoderFactory.get();

  private static ReadRowsResponse createResponse(
      Schema schema,
      Collection<GenericRecord> genericRecords,
      double progressAtResponseStart,
      double progressAtResponseEnd)
      throws Exception {
    GenericDatumWriter<GenericRecord> writer = new GenericDatumWriter<>(schema);
    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    Encoder binaryEncoder = ENCODER_FACTORY.binaryEncoder(outputStream, null);
    for (GenericRecord genericRecord : genericRecords) {
      writer.write(genericRecord, binaryEncoder);
    }

    binaryEncoder.flush();

    return ReadRowsResponse.newBuilder()
        .setAvroRows(
            AvroRows.newBuilder()
                .setSerializedBinaryRows(ByteString.copyFrom(outputStream.toByteArray()))
                .setRowCount(genericRecords.size()))
        .setRowCount(genericRecords.size())
        .setStats(
            StreamStats.newBuilder()
                .setProgress(
                    Progress.newBuilder()
                        .setAtResponseStart(progressAtResponseStart)
                        .setAtResponseEnd(progressAtResponseEnd)))
        .build();
  }

  private static final class ParseKeyValue
      implements SerializableFunction<SchemaAndRecord, KV<String, Long>> {
    @Override
    public KV<String, Long> apply(SchemaAndRecord input) {
      return KV.of(
          input.getRecord().get("name").toString(), (Long) input.getRecord().get("number"));
    }
  }

  @Test
  @ProjectOverride
  public void testQuerySourceInitialSplitWithBigQueryProject_EmptyResult() throws Exception {

    TableReference sourceTableRef =
        BigQueryHelpers.parseTableSpec("bigquery-project-id: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(TABLE_SCHEMA).setNumBytes(0L);

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable);

    fakeJobService.expectDryRunQuery(
        options.getBigQueryProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(1024L * 1024L)
                    .setReferencedTables(ImmutableList.of(sourceTableRef))));

    String stepUuid = "testStepUuid";

    TableReference tempTableReference =
        createTempTableReference(
            options.getBigQueryProject(),
            BigQueryResourceNaming.createJobIdPrefix(options.getJobName(), stepUuid, JobType.QUERY),
            Optional.empty());

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/" + options.getBigQueryProject())
            .setReadSession(
                ReadSession.newBuilder()
                    .setTable(BigQueryHelpers.toTableResourceName(tempTableReference))
                    .setDataFormat(DataFormat.AVRO))
            .setMaxStreamCount(10)
            .build();

    ReadSession emptyReadSession = ReadSession.newBuilder().build();
    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.createReadSession(expectedRequest)).thenReturn(emptyReadSession);

    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            stepUuid,
            ValueProvider.StaticValueProvider.of(encodedQuery),
            /* flattenResults = */ true,
            /* useLegacySql = */ true,
            /* priority = */ QueryPriority.BATCH,
            /* location = */ null,
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            DataFormat.AVRO,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withJobService(fakeJobService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = querySource.split(1024L, options);
    assertTrue(sources.isEmpty());
  }

  @Test
  public void testQuerySourceInitialSplit_EmptyResult() throws Exception {

    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(TABLE_SCHEMA).setNumBytes(0L);

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable);

    fakeJobService.expectDryRunQuery(
        options.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(1024L * 1024L)
                    .setReferencedTables(ImmutableList.of(sourceTableRef))));

    String stepUuid = "testStepUuid";

    TableReference tempTableReference =
        createTempTableReference(
            options.getProject(),
            BigQueryResourceNaming.createJobIdPrefix(options.getJobName(), stepUuid, JobType.QUERY),
            Optional.empty());

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/" + options.getProject())
            .setReadSession(
                ReadSession.newBuilder()
                    .setTable(BigQueryHelpers.toTableResourceName(tempTableReference)))
            .setMaxStreamCount(10)
            .build();

    ReadSession emptyReadSession = ReadSession.newBuilder().build();
    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.createReadSession(expectedRequest)).thenReturn(emptyReadSession);

    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            stepUuid,
            ValueProvider.StaticValueProvider.of(encodedQuery),
            /* flattenResults = */ true,
            /* useLegacySql = */ true,
            /* priority = */ QueryPriority.BATCH,
            /* location = */ null,
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withJobService(fakeJobService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = querySource.split(1024L, options);
    assertTrue(sources.isEmpty());
  }

  @Test
  public void testQuerySourceCreateReader() throws Exception {
    BigQueryStorageQuerySource<TableRow> querySource =
        BigQueryStorageQuerySource.create(
            /* stepUuid = */ "testStepUuid",
            ValueProvider.StaticValueProvider.of("SELECT * FROM `dataset.table`"),
            /* flattenResults = */ false,
            /* useLegacySql = */ false,
            /* priority = */ QueryPriority.INTERACTIVE,
            /* location = */ "asia-northeast1",
            /* queryTempDataset = */ null,
            /* kmsKey = */ null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            fakeBigQueryServices);

    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("BigQuery storage source must be split before reading");
    querySource.createReader(options);
  }

  @Test
  public void testReadFromBigQueryIO() throws Exception {
    doReadFromBigQueryIO(false);
  }

  @Test
  public void testReadFromBigQueryIOWithTemplateCompatibility() throws Exception {
    doReadFromBigQueryIO(true);
  }

  private void doReadFromBigQueryIO(boolean templateCompatibility) throws Exception {

    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(TABLE_SCHEMA).setNumBytes(0L);

    String encodedQuery = FakeBigQueryServices.encodeQueryResult(queryResultTable);

    fakeJobService.expectDryRunQuery(
        options.getProject(),
        encodedQuery,
        new JobStatistics()
            .setQuery(
                new JobStatistics2()
                    .setTotalBytesProcessed(1024L * 1024L)
                    .setReferencedTables(ImmutableList.of(sourceTableRef))));

    ReadSession readSession =
        ReadSession.newBuilder()
            .setName("readSessionName")
            .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
            .addStreams(ReadStream.newBuilder().setName("streamName"))
            .setDataFormat(DataFormat.AVRO)
            .build();

    ReadRowsRequest expectedReadRowsRequest =
        ReadRowsRequest.newBuilder().setReadStream("streamName").build();

    List<GenericRecord> records =
        Lists.newArrayList(
            createRecord("A", 1, AVRO_SCHEMA),
            createRecord("B", 2, AVRO_SCHEMA),
            createRecord("C", 3, AVRO_SCHEMA),
            createRecord("D", 4, AVRO_SCHEMA));

    List<ReadRowsResponse> readRowsResponses =
        Lists.newArrayList(
            createResponse(AVRO_SCHEMA, records.subList(0, 2), 0.0, 0.500),
            createResponse(AVRO_SCHEMA, records.subList(2, 4), 0.5, 0.875));

    //
    // Note that since the temporary table name is generated by the pipeline, we can't match the
    // expected create read session request exactly. For now, match against any appropriately typed
    // proto object.
    //

    StorageClient fakeStorageClient = mock(StorageClient.class, withSettings().serializable());
    when(fakeStorageClient.createReadSession(any())).thenReturn(readSession);
    when(fakeStorageClient.readRows(expectedReadRowsRequest, ""))
        .thenReturn(new FakeBigQueryServerStream<>(readRowsResponses));

    BigQueryIO.TypedRead<KV<String, Long>> typedRead =
        BigQueryIO.read(new ParseKeyValue())
            .fromQuery(encodedQuery)
            .withMethod(Method.DIRECT_READ)
            .withTestServices(
                new FakeBigQueryServices()
                    .withDatasetService(fakeDatasetService)
                    .withJobService(fakeJobService)
                    .withStorageClient(fakeStorageClient));

    if (templateCompatibility) {
      typedRead = typedRead.withTemplateCompatibility();
    }

    PCollection<KV<String, Long>> output = p.apply(typedRead);

    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(KV.of("A", 1L), KV.of("B", 2L), KV.of("C", 3L), KV.of("D", 4L)));

    p.run();
  }
}