/*
 * 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.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.hasItem;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.withSettings;

import com.google.api.gax.grpc.GrpcStatusCode;
import com.google.api.gax.rpc.FailedPreconditionException;
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.cloud.bigquery.storage.v1beta1.AvroProto.AvroRows;
import com.google.cloud.bigquery.storage.v1beta1.AvroProto.AvroSchema;
import com.google.cloud.bigquery.storage.v1beta1.ReadOptions.TableReadOptions;
import com.google.cloud.bigquery.storage.v1beta1.Storage.CreateReadSessionRequest;
import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsRequest;
import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsResponse;
import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession;
import com.google.cloud.bigquery.storage.v1beta1.Storage.SplitReadStreamRequest;
import com.google.cloud.bigquery.storage.v1beta1.Storage.SplitReadStreamResponse;
import com.google.cloud.bigquery.storage.v1beta1.Storage.Stream;
import com.google.cloud.bigquery.storage.v1beta1.Storage.StreamPosition;
import com.google.cloud.bigquery.storage.v1beta1.Storage.StreamStatus;
import com.google.protobuf.ByteString;
import com.google.protobuf.UnknownFieldSet;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.StatusRuntimeException;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Set;
import org.apache.avro.Schema;
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.BoundedSource.BoundedReader;
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.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.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
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.transforms.display.DisplayDataEvaluator;
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 BigQueryIOStorageReadTest {

  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(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();
                  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();

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

  @Test
  public void testBuildTableBasedSource() {
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table");
    checkTypedReadTableObject(typedRead, "foo.com:project", "dataset", "table");
    assertTrue(typedRead.getValidate());
  }

  @Test
  public void testBuildTableBasedSourceWithoutValidation() {
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .withoutValidation();
    checkTypedReadTableObject(typedRead, "foo.com:project", "dataset", "table");
    assertFalse(typedRead.getValidate());
  }

  @Test
  public void testBuildTableBasedSourceWithDefaultProject() {
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("myDataset.myTable");
    checkTypedReadTableObject(typedRead, null, "myDataset", "myTable");
  }

  @Test
  public void testBuildTableBasedSourceWithReadOptions() {
    TableReadOptions readOptions =
        TableReadOptions.newBuilder()
            .addSelectedFields("field1")
            .addSelectedFields("field2")
            .setRowRestriction("int_field > 5")
            .build();
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .withReadOptions(readOptions);
    checkTypedReadTableObject(typedRead, "foo.com:project", "dataset", "table");
    assertEquals(typedRead.getReadOptions(), readOptions);
  }

  @Test
  public void testBuildTableBasedSourceWithTableReference() {
    TableReference tableReference =
        new TableReference()
            .setProjectId("foo.com:project")
            .setDatasetId("dataset")
            .setTableId("table");
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from(tableReference);
    checkTypedReadTableObject(typedRead, "foo.com:project", "dataset", "table");
  }

  private void checkTypedReadTableObject(
      TypedRead typedRead, String project, String dataset, String table) {
    assertEquals(project, typedRead.getTable().getProjectId());
    assertEquals(dataset, typedRead.getTable().getDatasetId());
    assertEquals(table, typedRead.getTable().getTableId());
    assertNull(typedRead.getQuery());
    assertEquals(Method.DIRECT_READ, typedRead.getMethod());
  }

  @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(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .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(
        "ReadMyTable",
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .usingStandardSql());
    p.run();
  }

  @Test
  public void testBuildSourceWithReadOptionsAndSelectedFields() {
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage("withReadOptions() already called");
    p.apply(
        "ReadMyTable",
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .withReadOptions(TableReadOptions.newBuilder().build())
            .withSelectedFields(Lists.newArrayList("field1")));
  }

  @Test
  public void testBuildSourceWithReadOptionsAndRowRestriction() {
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage("withReadOptions() already called");
    p.apply(
        "ReadMyTable",
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .withReadOptions(TableReadOptions.newBuilder().build())
            .withRowRestriction("field > 1"));
  }

  @Test
  public void testDisplayData() {
    String tableSpec = "foo.com:project:dataset.table";
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from(tableSpec);
    DisplayData displayData = DisplayData.from(typedRead);
    assertThat(displayData, hasDisplayItem("table", tableSpec));
  }

  @Test
  public void testEvaluatedDisplayData() {
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    BigQueryIO.TypedRead<TableRow> typedRead =
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table");
    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(typedRead);
    assertThat(displayData, hasItem(hasDisplayItem("table")));
  }

  @Test
  public void testName() {
    assertEquals(
        "BigQueryIO.TypedRead",
        BigQueryIO.read(new TableRowParser())
            .withCoder(TableRowJsonCoder.of())
            .withMethod(Method.DIRECT_READ)
            .from("foo.com:project:dataset.table")
            .getName());
  }

  @Test
  public void testCoderInference() {
    // Lambdas erase too much type information -- use an anonymous class here.
    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 testTableSourceEstimatedSize() throws Exception {
    doTableSourceEstimatedSizeTest(false);
  }

  @Test
  public void testTableSourceEstimatedSize_IgnoresStreamingBuffer() throws Exception {
    doTableSourceEstimatedSizeTest(true);
  }

  private void doTableSourceEstimatedSizeTest(boolean useStreamingBuffer) throws Exception {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");
    Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
    if (useStreamingBuffer) {
      table.setStreamingBuffer(new Streamingbuffer().setEstimatedBytes(BigInteger.TEN));
    }

    fakeDatasetService.createTable(table);

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(tableRef),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withDatasetService(fakeDatasetService));

    assertEquals(100, tableSource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testTableSourceEstimatedSize_WithDefaultProject() throws Exception {
    fakeDatasetService.createDataset("project-id", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset.table");
    Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
    fakeDatasetService.createTable(table);

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withDatasetService(fakeDatasetService));

    assertEquals(100, tableSource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testTableSourceInitialSplit() throws Exception {
    doTableSourceInitialSplitTest(1024L, 1024);
  }

  @Test
  public void testTableSourceInitialSplit_MinSplitCount() throws Exception {
    doTableSourceInitialSplitTest(1024L * 1024L, 10);
  }

  @Test
  public void testTableSourceInitialSplit_MaxSplitCount() throws Exception {
    doTableSourceInitialSplitTest(10L, 10_000);
  }

  private void doTableSourceInitialSplitTest(long bundleSize, int streamCount) throws Exception {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");

    Table table =
        new Table()
            .setTableReference(tableRef)
            .setNumBytes(1024L * 1024L)
            .setSchema(new TableSchema());

    fakeDatasetService.createTable(table);

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(streamCount)
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

    ReadSession.Builder builder = ReadSession.newBuilder();
    for (int i = 0; i < streamCount; i++) {
      builder.addStreams(Stream.newBuilder().setName("stream-" + i));
    }

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

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(tableRef),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = tableSource.split(bundleSize, options);
    assertEquals(streamCount, sources.size());
  }

  @Test
  public void testTableSourceInitialSplit_WithTableReadOptions() throws Throwable {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");

    Table table =
        new Table()
            .setTableReference(tableRef)
            .setNumBytes(100L)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER"))));

    fakeDatasetService.createTable(table);

    TableReadOptions readOptions =
        TableReadOptions.newBuilder()
            .addSelectedFields("name")
            .addSelectedFields("number")
            .setRowRestriction("number > 5")
            .build();

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(10)
            .setReadOptions(readOptions)
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

    ReadSession.Builder builder = ReadSession.newBuilder();
    for (int i = 0; i < 10; i++) {
      builder.addStreams(Stream.newBuilder().setName("stream-" + i));
    }

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

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(tableRef),
            readOptions,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = tableSource.split(10L, options);
    assertEquals(10L, sources.size());
  }

  @Test
  public void testTableSourceInitialSplit_WithSelectedFieldsAndRowRestriction() throws Exception {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");

    Table table =
        new Table()
            .setTableReference(tableRef)
            .setNumBytes(100L)
            .setSchema(
                new TableSchema()
                    .setFields(
                        ImmutableList.of(
                            new TableFieldSchema().setName("name").setType("STRING"),
                            new TableFieldSchema().setName("number").setType("INTEGER"))));

    fakeDatasetService.createTable(table);

    TableReadOptions readOptions =
        TableReadOptions.newBuilder()
            .addSelectedFields("name")
            .addSelectedFields("number")
            .setRowRestriction("number > 5")
            .build();

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(10)
            .setReadOptions(readOptions)
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

    ReadSession.Builder builder = ReadSession.newBuilder();
    for (int i = 0; i < 10; i++) {
      builder.addStreams(Stream.newBuilder().setName("stream-" + i));
    }

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

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(tableRef),
            null,
            StaticValueProvider.of(Lists.newArrayList("name", "number")),
            StaticValueProvider.of("number > 5"),
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withStorageClient(fakeStorageClient));

    List<? extends BoundedSource<TableRow>> sources = tableSource.split(10L, options);
    assertEquals(10L, sources.size());
  }

  @Test
  public void testTableSourceInitialSplit_WithDefaultProject() throws Exception {
    fakeDatasetService.createDataset("project-id", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset.table");

    Table table =
        new Table()
            .setTableReference(tableRef)
            .setNumBytes(1024L * 1024L)
            .setSchema(new TableSchema());

    fakeDatasetService.createTable(table);

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(1024)
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

    ReadSession.Builder builder = ReadSession.newBuilder();
    for (int i = 0; i < 50; i++) {
      builder.addStreams(Stream.newBuilder().setName("stream-" + i));
    }

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

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withStorageClient(fakeStorageClient));

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

  @Test
  public void testTableSourceInitialSplit_EmptyTable() throws Exception {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");

    Table table =
        new Table()
            .setTableReference(tableRef)
            .setNumBytes(1024L * 1024L)
            .setSchema(new TableSchema());

    fakeDatasetService.createTable(table);

    CreateReadSessionRequest expectedRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(1024)
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

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

    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(tableRef),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices()
                .withDatasetService(fakeDatasetService)
                .withStorageClient(fakeStorageClient));

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

  @Test
  public void testTableSourceCreateReader() throws Exception {
    BigQueryStorageTableSource<TableRow> tableSource =
        BigQueryStorageTableSource.create(
            ValueProvider.StaticValueProvider.of(
                BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table")),
            null,
            null,
            null,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withDatasetService(fakeDatasetService));

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

  private static final String AVRO_SCHEMA_STRING =
      "{\"namespace\": \"example.avro\",\n"
          + " \"type\": \"record\",\n"
          + " \"name\": \"RowRecord\",\n"
          + " \"fields\": [\n"
          + "     {\"name\": \"name\", \"type\": \"string\"},\n"
          + "     {\"name\": \"number\", \"type\": \"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").setMode("REQUIRED"),
                  new TableFieldSchema().setName("number").setType("INTEGER").setMode("REQUIRED")));

  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 fractionConsumed)
      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()))
        .setStatus(
            StreamStatus.newBuilder()
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFractionConsumed().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(
                                    java.lang.Float.floatToIntBits((float) fractionConsumed))
                                .build())
                        .build()))
        .build();
  }

  @Test
  public void testStreamSourceEstimatedSizeBytes() throws Exception {

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            ReadSession.getDefaultInstance(),
            Stream.getDefaultInstance(),
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices());

    assertEquals(0, streamSource.getEstimatedSizeBytes(options));
  }

  @Test
  public void testStreamSourceSplit() throws Exception {

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            ReadSession.getDefaultInstance(),
            Stream.getDefaultInstance(),
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices());

    assertThat(streamSource.split(0, options), containsInAnyOrder(streamSource));
  }

  @Test
  public void testReadFromStreamSource() throws Exception {

    ReadSession readSession =
        ReadSession.newBuilder()
            .setName("readSession")
            .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
            .build();

    Stream stream = Stream.newBuilder().setName("stream").build();

    ReadRowsRequest expectedRequest =
        ReadRowsRequest.newBuilder()
            .setReadPosition(StreamPosition.newBuilder().setStream(stream))
            .build();

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

    List<ReadRowsResponse> responses =
        Lists.newArrayList(
            createResponse(AVRO_SCHEMA, records.subList(0, 2), 0.50),
            createResponse(AVRO_SCHEMA, records.subList(2, 3), 0.75));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(expectedRequest))
        .thenReturn(new FakeBigQueryServerStream<>(responses));

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            readSession,
            stream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    List<TableRow> rows = new ArrayList<>();
    BoundedReader<TableRow> reader = streamSource.createReader(options);
    for (boolean hasNext = reader.start(); hasNext; hasNext = reader.advance()) {
      rows.add(reader.getCurrent());
    }

    System.out.println("Rows: " + rows);

    assertEquals(3, rows.size());
  }

  @Test
  public void testFractionConsumed() throws Exception {
    ReadSession readSession =
        ReadSession.newBuilder()
            .setName("readSession")
            .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
            .build();

    Stream stream = Stream.newBuilder().setName("stream").build();

    ReadRowsRequest expectedRequest =
        ReadRowsRequest.newBuilder()
            .setReadPosition(StreamPosition.newBuilder().setStream(stream))
            .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),
            createRecord("E", 5, AVRO_SCHEMA),
            createRecord("F", 6, AVRO_SCHEMA),
            createRecord("G", 7, AVRO_SCHEMA));

    List<ReadRowsResponse> responses =
        Lists.newArrayList(
            // N.B.: All floating point numbers used in this test can be represented without
            // a loss of precision.
            createResponse(AVRO_SCHEMA, records.subList(0, 2), 0.250),
            createResponse(AVRO_SCHEMA, records.subList(2, 4), 0.500),
            createResponse(AVRO_SCHEMA, records.subList(4, 7), 0.875));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(expectedRequest))
        .thenReturn(new FakeBigQueryServerStream<>(responses));

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            readSession,
            stream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    List<TableRow> rows = new ArrayList<>();
    BoundedReader<TableRow> reader = streamSource.createReader(options);

    // Before call to BoundedReader#start, fraction consumed must be zero.
    assertEquals(Double.valueOf(0.000), reader.getFractionConsumed());

    assertTrue(reader.start()); // Reads A.
    assertEquals(Double.valueOf(0.125), reader.getFractionConsumed());
    assertTrue(reader.advance()); // Reads B.
    assertEquals(Double.valueOf(0.250), reader.getFractionConsumed());

    assertTrue(reader.advance()); // Reads C.
    assertEquals(Double.valueOf(0.375), reader.getFractionConsumed());
    assertTrue(reader.advance()); // Reads D.
    assertEquals(Double.valueOf(0.500), reader.getFractionConsumed());

    assertTrue(reader.advance()); // Reads E.
    assertEquals(Double.valueOf(0.625), reader.getFractionConsumed());
    assertTrue(reader.advance()); // Reads F.
    assertEquals(Double.valueOf(0.750), reader.getFractionConsumed());
    assertTrue(reader.advance()); // Reads G.
    assertEquals(Double.valueOf(0.875), reader.getFractionConsumed());

    assertFalse(reader.advance()); // Reaches the end.

    // We are done with the stream, so we should report 100% consumption.
    assertEquals(Double.valueOf(1.00), reader.getFractionConsumed());
  }

  @Test
  public void testFractionConsumedWithSplit() throws Exception {
    ReadSession readSession =
        ReadSession.newBuilder()
            .setName("readSession")
            .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
            .build();

    Stream parentStream = Stream.newBuilder().setName("stream").build();

    ReadRowsRequest expectedRequest =
        ReadRowsRequest.newBuilder()
            .setReadPosition(StreamPosition.newBuilder().setStream(parentStream))
            .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),
            createRecord("E", 5, AVRO_SCHEMA),
            createRecord("F", 6, AVRO_SCHEMA),
            createRecord("G", 7, AVRO_SCHEMA));

    List<ReadRowsResponse> parentResponses =
        Lists.newArrayList(
            // N.B.: All floating point numbers used in this test can be represented without
            // a loss of precision.
            createResponse(AVRO_SCHEMA, records.subList(0, 2), 0.250),
            createResponse(AVRO_SCHEMA, records.subList(2, 4), 0.500),
            createResponse(AVRO_SCHEMA, records.subList(4, 7), 0.875));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(expectedRequest))
        .thenReturn(new FakeBigQueryServerStream<>(parentResponses));

    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(parentStream)
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.5f))
                                .build())
                        .build())
                .build()))
        .thenReturn(
            SplitReadStreamResponse.newBuilder()
                .setPrimaryStream(Stream.newBuilder().setName("primary"))
                .setRemainderStream(Stream.newBuilder().setName("residual"))
                .build());

    List<ReadRowsResponse> primaryResponses =
        Lists.newArrayList(
            createResponse(AVRO_SCHEMA, records.subList(1, 3), 0.500),
            createResponse(AVRO_SCHEMA, records.subList(3, 4), 0.875));

    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(
                    StreamPosition.newBuilder()
                        .setStream(Stream.newBuilder().setName("primary"))
                        .setOffset(1))
                .build()))
        .thenReturn(new FakeBigQueryServerStream<>(primaryResponses));

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            readSession,
            parentStream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    List<TableRow> rows = new ArrayList<>();
    BoundedReader<TableRow> reader = streamSource.createReader(options);

    // Before call to BoundedReader#start, fraction consumed must be zero.
    assertEquals(Double.valueOf(0.0000), reader.getFractionConsumed());

    assertTrue(reader.start()); // Reads A.
    assertEquals(Double.valueOf(0.1250), reader.getFractionConsumed());

    reader.splitAtFraction(0.5f);
    assertEquals(Double.valueOf(0.1250), reader.getFractionConsumed());

    assertTrue(reader.advance()); // Reads B.

    // Once the split has completed but no new rows have been read, the consumed value is at the
    // last calculated value of 0.125. For the first response of the primary stream, the progress
    // report interpolation is done between the progress before split and the progress from the
    // first response of the primary stream. In this case, the value is:
    //
    //   0.125 + (0.5 - 0.125) * 1.0 / 2
    //
    assertEquals(Double.valueOf(0.3125), reader.getFractionConsumed());

    assertTrue(reader.advance()); // Reads C.
    assertEquals(Double.valueOf(0.5000), reader.getFractionConsumed());

    assertTrue(reader.advance()); // Reads D.
    assertEquals(Double.valueOf(0.8750), reader.getFractionConsumed());

    assertFalse(reader.advance());
    assertEquals(Double.valueOf(1.0000), reader.getFractionConsumed());
  }

  @Test
  public void testStreamSourceSplitAtFractionSucceeds() throws Exception {
    Stream parentStream = Stream.newBuilder().setName("parent").build();

    List<ReadRowsResponse> parentResponses =
        Lists.newArrayList(
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("A", 1, AVRO_SCHEMA), createRecord("B", 2, AVRO_SCHEMA)),
                0.25),
            createResponse(
                AVRO_SCHEMA, Lists.newArrayList(createRecord("C", 3, AVRO_SCHEMA)), 0.50),
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("D", 4, AVRO_SCHEMA), createRecord("E", 5, AVRO_SCHEMA)),
                0.75));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(parentStream))
                .build()))
        .thenReturn(new FakeBigQueryServerStream<>(parentResponses));

    // Mocks the split call.
    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(parentStream)
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.5f))
                                .build())
                        .build())
                .build()))
        .thenReturn(
            SplitReadStreamResponse.newBuilder()
                .setPrimaryStream(Stream.newBuilder().setName("primary"))
                .setRemainderStream(Stream.newBuilder().setName("residual"))
                .build());

    // Mocks the ReadRows calls expected on the primary and residual streams.
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(
                    StreamPosition.newBuilder()
                        .setStream(Stream.newBuilder().setName("primary"))
                        // This test will read rows 0 and 1 from the parent before calling split,
                        // so we expect the primary read to start at offset 2.
                        .setOffset(2))
                .build()))
        .thenReturn(new FakeBigQueryServerStream<>(parentResponses.subList(1, 2)));
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(
                    StreamPosition.newBuilder()
                        .setStream(Stream.newBuilder().setName("residual"))
                        .setOffset(0))
                .build()))
        .thenReturn(
            new FakeBigQueryServerStream<>(parentResponses.subList(2, parentResponses.size())));

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            ReadSession.newBuilder()
                .setName("readSession")
                .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
                .build(),
            parentStream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    // Read a few records from the parent stream and ensure that records are returned in the
    // prescribed order.
    BoundedReader<TableRow> parent = streamSource.createReader(options);
    assertTrue(parent.start());
    assertEquals("A", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("B", parent.getCurrent().get("name"));

    // Now split the stream, and ensure that the "parent" reader has been replaced with the
    // primary stream and that the returned source points to the residual stream.
    BoundedReader<TableRow> primary = parent;
    BoundedSource<TableRow> residualSource = parent.splitAtFraction(0.5);
    assertNotNull(residualSource);
    BoundedReader<TableRow> residual = residualSource.createReader(options);

    assertTrue(primary.advance());
    assertEquals("C", primary.getCurrent().get("name"));
    assertFalse(primary.advance());

    assertTrue(residual.start());
    assertEquals("D", residual.getCurrent().get("name"));
    assertTrue(residual.advance());
    assertEquals("E", residual.getCurrent().get("name"));
    assertFalse(residual.advance());
  }

  @Test
  public void testStreamSourceSplitAtFractionRepeated() throws Exception {
    List<Stream> streams =
        Lists.newArrayList(
            Stream.newBuilder().setName("stream1").build(),
            Stream.newBuilder().setName("stream2").build(),
            Stream.newBuilder().setName("stream3").build());

    StorageClient fakeStorageClient = mock(StorageClient.class);

    // Mock the initial ReadRows call.
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(streams.get(0)))
                .build()))
        .thenReturn(
            new FakeBigQueryServerStream<>(
                Lists.newArrayList(
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("A", 1, AVRO_SCHEMA), createRecord("B", 2, AVRO_SCHEMA)),
                        0.25),
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("C", 3, AVRO_SCHEMA), createRecord("D", 4, AVRO_SCHEMA)),
                        0.50),
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("E", 5, AVRO_SCHEMA), createRecord("F", 6, AVRO_SCHEMA)),
                        0.75))));

    // Mock the first SplitReadStream call.
    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(streams.get(0))
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.83f))
                                .build())
                        .build())
                .build()))
        .thenReturn(
            SplitReadStreamResponse.newBuilder()
                .setPrimaryStream(streams.get(1))
                .setRemainderStream(Stream.newBuilder().setName("ignored"))
                .build());

    // Mock the second ReadRows call.
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(streams.get(1)).setOffset(1))
                .build()))
        .thenReturn(
            new FakeBigQueryServerStream<>(
                Lists.newArrayList(
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("B", 2, AVRO_SCHEMA), createRecord("C", 3, AVRO_SCHEMA)),
                        0.50),
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("D", 4, AVRO_SCHEMA), createRecord("E", 5, AVRO_SCHEMA)),
                        0.75))));

    // Mock the second SplitReadStream call.
    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(streams.get(1))
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.75f))
                                .build())
                        .build())
                .build()))
        .thenReturn(
            SplitReadStreamResponse.newBuilder()
                .setPrimaryStream(streams.get(2))
                .setRemainderStream(Stream.newBuilder().setName("ignored"))
                .build());

    // Mock the third ReadRows call.
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(streams.get(2)).setOffset(2))
                .build()))
        .thenReturn(
            new FakeBigQueryServerStream<>(
                Lists.newArrayList(
                    createResponse(
                        AVRO_SCHEMA,
                        Lists.newArrayList(
                            createRecord("C", 3, AVRO_SCHEMA), createRecord("D", 4, AVRO_SCHEMA)),
                        0.90))));

    BoundedSource<TableRow> source =
        BigQueryStorageStreamSource.create(
            ReadSession.newBuilder()
                .setName("readSession")
                .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
                .build(),
            streams.get(0),
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    BoundedReader<TableRow> reader = source.createReader(options);
    assertTrue(reader.start());
    assertEquals("A", reader.getCurrent().get("name"));

    BoundedSource<TableRow> residualSource = reader.splitAtFraction(0.83f);
    assertNotNull(residualSource);
    assertEquals("A", reader.getCurrent().get("name"));

    assertTrue(reader.advance());
    assertEquals("B", reader.getCurrent().get("name"));

    residualSource = reader.splitAtFraction(0.75f);
    assertNotNull(residualSource);
    assertEquals("B", reader.getCurrent().get("name"));

    assertTrue(reader.advance());
    assertEquals("C", reader.getCurrent().get("name"));
    assertTrue(reader.advance());
    assertEquals("D", reader.getCurrent().get("name"));
    assertFalse(reader.advance());
  }

  @Test
  public void testStreamSourceSplitAtFractionFailsWhenSplitIsNotPossible() throws Exception {
    Stream parentStream = Stream.newBuilder().setName("parent").build();

    List<ReadRowsResponse> parentResponses =
        Lists.newArrayList(
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("A", 1, AVRO_SCHEMA), createRecord("B", 2, AVRO_SCHEMA)),
                0.25),
            createResponse(
                AVRO_SCHEMA, Lists.newArrayList(createRecord("C", 3, AVRO_SCHEMA)), 0.50),
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("D", 4, AVRO_SCHEMA), createRecord("E", 5, AVRO_SCHEMA)),
                0.75));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(parentStream))
                .build()))
        .thenReturn(new FakeBigQueryServerStream<>(parentResponses));

    // Mocks the split call. A response without a primary_stream and remainder_stream means
    // that the split is not possible.
    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(parentStream)
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.5f))
                                .build())
                        .build())
                .build()))
        .thenReturn(SplitReadStreamResponse.getDefaultInstance());

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            ReadSession.newBuilder()
                .setName("readSession")
                .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
                .build(),
            parentStream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    // Read a few records from the parent stream and ensure that records are returned in the
    // prescribed order.
    BoundedReader<TableRow> parent = streamSource.createReader(options);
    assertTrue(parent.start());
    assertEquals("A", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("B", parent.getCurrent().get("name"));

    assertNull(parent.splitAtFraction(0.5));

    // Verify that the parent source still works okay even after an unsuccessful split attempt.
    assertTrue(parent.advance());
    assertEquals("C", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("D", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("E", parent.getCurrent().get("name"));
    assertFalse(parent.advance());
  }

  @Test
  public void testStreamSourceSplitAtFractionFailsWhenParentIsPastSplitPoint() throws Exception {
    Stream parentStream = Stream.newBuilder().setName("parent").build();

    List<ReadRowsResponse> parentResponses =
        Lists.newArrayList(
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("A", 1, AVRO_SCHEMA), createRecord("B", 2, AVRO_SCHEMA)),
                0.25),
            createResponse(
                AVRO_SCHEMA, Lists.newArrayList(createRecord("C", 3, AVRO_SCHEMA)), 0.50),
            createResponse(
                AVRO_SCHEMA,
                Lists.newArrayList(
                    createRecord("D", 4, AVRO_SCHEMA), createRecord("E", 5, AVRO_SCHEMA)),
                0.75));

    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(StreamPosition.newBuilder().setStream(parentStream))
                .build()))
        .thenReturn(new FakeBigQueryServerStream<>(parentResponses));

    // Mocks the split call. A response without a primary_stream and remainder_stream means
    // that the split is not possible.
    // Mocks the split call.
    when(fakeStorageClient.splitReadStream(
            SplitReadStreamRequest.newBuilder()
                .setOriginalStream(parentStream)
                // TODO(aryann): Once we rebuild the generated client code, we should change this to
                // use setFraction().
                .setUnknownFields(
                    UnknownFieldSet.newBuilder()
                        .addField(
                            2,
                            UnknownFieldSet.Field.newBuilder()
                                .addFixed32(java.lang.Float.floatToIntBits(0.5f))
                                .build())
                        .build())
                .build()))
        .thenReturn(
            SplitReadStreamResponse.newBuilder()
                .setPrimaryStream(Stream.newBuilder().setName("primary"))
                .setRemainderStream(Stream.newBuilder().setName("residual"))
                .build());

    // Mocks the ReadRows calls expected on the primary and residual streams.
    when(fakeStorageClient.readRows(
            ReadRowsRequest.newBuilder()
                .setReadPosition(
                    StreamPosition.newBuilder()
                        .setStream(Stream.newBuilder().setName("primary"))
                        // This test will read rows 0 and 1 from the parent before calling split,
                        // so we expect the primary read to start at offset 2.
                        .setOffset(2))
                .build()))
        .thenThrow(
            new FailedPreconditionException(
                "Given row offset is invalid for stream.",
                new StatusRuntimeException(Status.FAILED_PRECONDITION),
                GrpcStatusCode.of(Code.FAILED_PRECONDITION),
                /* retryable = */ false));

    BigQueryStorageStreamSource<TableRow> streamSource =
        BigQueryStorageStreamSource.create(
            ReadSession.newBuilder()
                .setName("readSession")
                .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
                .build(),
            parentStream,
            TABLE_SCHEMA,
            new TableRowParser(),
            TableRowJsonCoder.of(),
            new FakeBigQueryServices().withStorageClient(fakeStorageClient));

    // Read a few records from the parent stream and ensure that records are returned in the
    // prescribed order.
    BoundedReader<TableRow> parent = streamSource.createReader(options);
    assertTrue(parent.start());
    assertEquals("A", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("B", parent.getCurrent().get("name"));

    assertNull(parent.splitAtFraction(0.5));

    // Verify that the parent source still works okay even after an unsuccessful split attempt.
    assertTrue(parent.advance());
    assertEquals("C", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("D", parent.getCurrent().get("name"));
    assertTrue(parent.advance());
    assertEquals("E", parent.getCurrent().get("name"));
    assertFalse(parent.advance());
  }

  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
  public void testReadFromBigQueryIO() throws Exception {
    fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");

    Table table =
        new Table().setTableReference(tableRef).setNumBytes(10L).setSchema(new TableSchema());

    fakeDatasetService.createTable(table);

    CreateReadSessionRequest expectedCreateReadSessionRequest =
        CreateReadSessionRequest.newBuilder()
            .setParent("projects/project-id")
            .setTableReference(BigQueryHelpers.toTableRefProto(tableRef))
            .setRequestedStreams(10)
            .setReadOptions(
                TableReadOptions.newBuilder().addSelectedFields("name").addSelectedFields("number"))
            // TODO(aryann): Once we rebuild the generated client code, we should change this to
            // use setShardingStrategy().
            .setUnknownFields(
                UnknownFieldSet.newBuilder()
                    .addField(7, UnknownFieldSet.Field.newBuilder().addVarint(2).build())
                    .build())
            .build();

    ReadSession readSession =
        ReadSession.newBuilder()
            .setName("readSessionName")
            .setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING))
            .addStreams(Stream.newBuilder().setName("streamName"))
            .build();

    ReadRowsRequest expectedReadRowsRequest =
        ReadRowsRequest.newBuilder()
            .setReadPosition(
                StreamPosition.newBuilder().setStream(Stream.newBuilder().setName("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.50),
            createResponse(AVRO_SCHEMA, records.subList(2, 4), 0.75));

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

    PCollection<KV<String, Long>> output =
        p.apply(
            BigQueryIO.read(new ParseKeyValue())
                .from("foo.com:project:dataset.table")
                .withMethod(Method.DIRECT_READ)
                .withSelectedFields(p.newProvider(Lists.newArrayList("name", "number")))
                .withTestServices(
                    new FakeBigQueryServices()
                        .withDatasetService(fakeDatasetService)
                        .withStorageClient(fakeStorageClient)));

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

    p.run();
  }
}