Add integration tests for snapshot read and resuming.

google-cloud-clients/google-cloud-bigquerystorage/pom.xml

@@ -41,6 +41,11 @@
             <artifactId>gax-grpc</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-bigquery</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.avro</groupId>
             <artifactId>avro</artifactId>

google-cloud-clients/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java

@@ -18,36 +18,67 @@
 
 import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.gax.rpc.ServerStream;
+import com.google.cloud.RetryOption;
 import com.google.cloud.ServiceOptions;
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.Field.Mode;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.JobInfo;
+import com.google.cloud.bigquery.JobInfo.WriteDisposition;
+import com.google.cloud.bigquery.LegacySQLTypeName;
+import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.StandardTableDefinition;
+import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient;
 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.DataFormat;
 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.Stream;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.StreamPosition;
+import com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableModifiers;
 import com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference;
+import com.google.cloud.bigquery.storage.v1beta1.it.SimpleRowReader.AvroRowConsumer;
+import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
+import com.google.common.base.Preconditions;
 import com.google.protobuf.TextFormat;
+import com.google.protobuf.Timestamp;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
 import java.util.logging.Logger;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.util.Utf8;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.threeten.bp.Duration;
+import org.threeten.bp.Instant;
 
 /** Integration tests for BigQuery Storage API. */
 public class ITBigQueryStorageTest {
 
   private static final Logger LOG = Logger.getLogger(ITBigQueryStorageTest.class.getName());
+  private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
+  private static final String DESCRIPTION = "BigQuery Storage Java client test dataset";
 
   private static BigQueryStorageClient client;
   private static String parentProjectId;
+  private static BigQuery bigquery;
 
   @BeforeClass
   public static void beforeClass() throws IOException {
@@ -58,13 +89,25 @@ public static void beforeClass() throws IOException {
         String.format(
             "%s tests running with parent project: %s",
             ITBigQueryStorageTest.class.getSimpleName(), parentProjectId));
+
+    RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+    bigquery = bigqueryHelper.getOptions().getService();
+    DatasetInfo datasetInfo =
+        DatasetInfo.newBuilder(/* datasetId = */ DATASET).setDescription(DESCRIPTION).build();
+    bigquery.create(datasetInfo);
+    LOG.info("Created test dataset: " + DATASET);
   }
 
   @AfterClass
   public static void afterClass() {
     if (client != null) {
       client.close();
     }
+
+    if (bigquery != null) {
+      RemoteBigQueryHelper.forceDelete(bigquery, DATASET);
+      LOG.info("Deleted test dataset: " + DATASET);
+    }
   }
 
   @Test
@@ -76,7 +119,11 @@ public void testSimpleRead() {
             .setTableId("shakespeare")
             .build();
 
-    ReadSession session = client.createReadSession(tableReference, parentProjectId, 1);
+    ReadSession session =
+        client.createReadSession(
+            /* tableReference = */ tableReference,
+            /* parent = */ parentProjectId,
+            /* requestedStreams = */ 1);
     assertEquals(
         String.format(
             "Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s",
@@ -104,6 +151,57 @@ public void testSimpleRead() {
     assertEquals(164_656, avroRowCount);
   }
 
+  @Test
+  public void testSimpleReadAndResume() {
+    TableReference tableReference =
+        TableReference.newBuilder()
+            .setProjectId("bigquery-public-data")
+            .setDatasetId("samples")
+            .setTableId("shakespeare")
+            .build();
+
+    ReadSession session =
+        client.createReadSession(
+            /* tableReference = */ tableReference,
+            /* parent = */ parentProjectId,
+            /* requestedStreams = */ 1);
+    assertEquals(
+        String.format(
+            "Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s",
+            TextFormat.shortDebugString(tableReference), session.toString()),
+        1,
+        session.getStreamsCount());
+
+    // We have to read some number of rows in order to be able to resume. More details:
+    // https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1beta1#google.cloud.bigquery.storage.v1beta1.ReadRowsRequest
+
+    long avroRowCount = SkipNumberOfRows(session.getStreams(0), /* rowsToSkip = */ 34_846);
+
+    StreamPosition readPosition =
+        StreamPosition.newBuilder()
+            .setStream(session.getStreams(0))
+            .setOffset(avroRowCount)
+            .build();
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadPosition(readPosition).build();
+
+    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
+
+    for (ReadRowsResponse response : stream) {
+      assertTrue(
+          String.format(
+              "Response is missing 'avro_rows'. Read %d rows so far. ReadRows response:%n%s",
+              avroRowCount, response.toString()),
+          response.hasAvroRows());
+      avroRowCount += response.getAvroRows().getRowCount();
+    }
+
+    // Verifies that the number of rows skipped and read equals to the total number of rows in the
+    // table.
+    assertEquals(164_656, avroRowCount);
+  }
+
   @Test
   public void testFilter() throws IOException {
     TableReference tableReference =
@@ -254,4 +352,174 @@ public void accept(GenericRecord record) {
 
     assertEquals(1_333, avroRowCount);
   }
+
+  @Test
+  public void testReadAtSnapshot() throws InterruptedException, IOException {
+    Field intFieldSchema =
+        Field.newBuilder("col", LegacySQLTypeName.INTEGER)
+            .setMode(Mode.REQUIRED)
+            .setDescription("IntegerDescription")
+            .build();
+    com.google.cloud.bigquery.Schema tableSchema =
+        com.google.cloud.bigquery.Schema.of(intFieldSchema);
+
+    TableId testTableId = TableId.of(/* dataset = */ DATASET, /* table = */ "test_read_snapshot");
+    bigquery.create(TableInfo.of(testTableId, StandardTableDefinition.of(tableSchema)));
+
+    TableReference tableReference =
+        TableReference.newBuilder()
+            .setTableId(testTableId.getTable())
+            .setDatasetId(DATASET)
+            .setProjectId(ServiceOptions.getDefaultProjectId())
+            .build();
+
+    RunQueryJobAndExpectSuccess(
+        /* destinationTableId = */ testTableId, /* query = */ "SELECT 1 AS col");
+    Instant firstSnapshot = Instant.now();
+
+    RunQueryJobAndExpectSuccess(
+        /* destinationTableId = */ testTableId, /* query = */ "SELECT 2 AS col");
+    Instant secondSnapshot = Instant.now();
+
+    final List<Long> rowsAfterFirstSnapshot = new ArrayList<>();
+    ProcessRowsAtSnapshot(
+        tableReference,
+        firstSnapshot,
+        new AvroRowConsumer() {
+          @Override
+          public void accept(GenericRecord record) {
+            rowsAfterFirstSnapshot.add((Long) record.get("col"));
+          }
+        });
+    assertEquals(Arrays.asList(1L), rowsAfterFirstSnapshot);
+
+    final List<Long> rowsAfterSecondSnapshot = new ArrayList<>();
+    ProcessRowsAtSnapshot(
+        tableReference,
+        secondSnapshot,
+        new AvroRowConsumer() {
+          @Override
+          public void accept(GenericRecord record) {
+            rowsAfterSecondSnapshot.add((Long) record.get("col"));
+          }
+        });
+    Collections.sort(rowsAfterSecondSnapshot);
+    assertEquals(Arrays.asList(1L, 2L), rowsAfterSecondSnapshot);
+  }
+
+  /**
+   * Skips the specified number of rows from the stream. If the stream does not have the desired
+   * rows to skip, it will skip all of them.
+   *
+   * @param stream
+   * @param rowsToSkip
+   * @return the number of requested rows to skip or the total rows read if stream had less rows.
+   */
+  private long SkipNumberOfRows(Stream stream, long rowsToSkip) {
+    StreamPosition readPosition = StreamPosition.newBuilder().setStream(stream).build();
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadPosition(readPosition).build();
+
+    long avroRowCount = 0;
+    ServerStream<ReadRowsResponse> serverStream = client.readRowsCallable().call(readRowsRequest);
+    Iterator<ReadRowsResponse> responseIterator = serverStream.iterator();
+
+    while (responseIterator.hasNext()) {
+      ReadRowsResponse response = responseIterator.next();
+      avroRowCount += response.getAvroRows().getRowCount();
+      if (avroRowCount >= rowsToSkip) {
+        return rowsToSkip;
+      }
+    }
+
+    return avroRowCount;
+  }
+
+  /**
+   * Reads all the rows from the specified tableReference that are added up to timestamp defined in
+   * snapshot. If snapshot is not provided, current time will be used.
+   *
+   * <p>For every row, the consumer is called for processing.
+   *
+   * @param tableReference
+   * @param snapshot
+   * @param consumer
+   * @throws IOException
+   */
+  private void ProcessRowsAtSnapshot(
+      TableReference tableReference, Instant snapshot, AvroRowConsumer consumer)
+      throws IOException {
+    Preconditions.checkNotNull(tableReference);
+    Preconditions.checkNotNull(consumer);
+
+    CreateReadSessionRequest.Builder createSessionRequestBuilder =
+        CreateReadSessionRequest.newBuilder()
+            .setParent(parentProjectId)
+            .setRequestedStreams(1)
+            .setTableReference(tableReference)
+            .setFormat(DataFormat.AVRO);
+
+    if (snapshot != null) {
+      Timestamp snapshotTimestamp =
+          Timestamp.newBuilder()
+              .setSeconds(snapshot.toEpochMilli() / 1_000)
+              .setNanos((int) ((snapshot.toEpochMilli() % 1000) * 1000000))
+              .build();
+      createSessionRequestBuilder.setTableModifiers(
+          TableModifiers.newBuilder().setSnapshotTime(snapshotTimestamp).build());
+    }
+
+    ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+    assertEquals(
+        String.format(
+            "Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s",
+            TextFormat.shortDebugString(tableReference), session.toString()),
+        1,
+        session.getStreamsCount());
+
+    StreamPosition readPosition =
+        StreamPosition.newBuilder().setStream(session.getStreams(0)).build();
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadPosition(readPosition).build();
+
+    SimpleRowReader reader =
+        new SimpleRowReader(new Schema.Parser().parse(session.getAvroSchema().getSchema()));
+
+    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
+    for (ReadRowsResponse response : stream) {
+      reader.processRows(response.getAvroRows(), consumer);
+    }
+  }
+
+  /**
+   * Runs a query job with WRITE_APPEND disposition to the destination table.
+   *
+   * @param destinationTableId
+   * @param query
+   * @throws InterruptedException
+   */
+  private void RunQueryJobAndExpectSuccess(TableId destinationTableId, String query)
+      throws InterruptedException {
+    QueryJobConfiguration configuration =
+        QueryJobConfiguration.newBuilder(query)
+            .setDestinationTable(destinationTableId)
+            .setUseQueryCache(false)
+            .setUseLegacySql(false)
+            .setWriteDisposition(WriteDisposition.WRITE_APPEND)
+            .build();
+
+    Job job = bigquery.create(JobInfo.of(configuration));
+    Job completedJob =
+        job.waitFor(
+            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
+            RetryOption.totalTimeout(Duration.ofMinutes(1)));
+
+    assertNotNull(completedJob);
+    assertNull(
+        /* message = */ "Received a job status that is not a success: "
+            + completedJob.getStatus().toString(),
+        /* object = */ completedJob.getStatus().getError());
+  }
 }