test: Add timestamp ITs for Read and Write API

google-cloud-bigquerystorage/pom.xml

@@ -232,6 +232,11 @@
       <artifactId>arrow-memory-core</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java-util</artifactId>
+      <scope>test</scope>
+    </dependency>
 
     <dependency>
       <groupId>io.grpc</groupId>

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageLongRunningTest.java

@@ -26,6 +26,7 @@
 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.it.util.BigQueryResource;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteNonQuotaRetryTest.java

@@ -38,6 +38,7 @@
 import com.google.cloud.bigquery.storage.v1.StreamWriter;
 import com.google.cloud.bigquery.storage.v1.TableName;
 import com.google.cloud.bigquery.storage.v1.WriteStream;
+import com.google.cloud.bigquery.storage.v1.it.util.WriteRetryTestUtil;
 import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
 import com.google.protobuf.DescriptorProtos.DescriptorProto;
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteQuotaRetryTest.java

@@ -26,6 +26,7 @@
 import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
 import com.google.cloud.bigquery.storage.v1.WriteStream;
+import com.google.cloud.bigquery.storage.v1.it.util.WriteRetryTestUtil;
 import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import java.io.IOException;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/BigQueryResource.java → google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/BigQueryResource.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.google.cloud.bigquery.storage.v1.it;
+package com.google.cloud.bigquery.storage.v1.it.util;
 
 /** Test helper class to generate BigQuery resource paths. */
 public class BigQueryResource {

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java

@@ -0,0 +1,180 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed 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 com.google.cloud.bigquery.storage.v1.it.util;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.gax.rpc.ServerStream;
+import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.cloud.bigquery.storage.v1.AppendRowsResponse;
+import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
+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.common.base.Preconditions;
+import com.google.protobuf.util.Timestamps;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericRecordBuilder;
+
+public class Helper {
+
+  public static final Long[] INPUT_TIMESTAMPS_MICROS =
+      new Long[] {
+        1735734896123456L, // 2025-01-01T12:34:56.123456Z
+        1580646896123456L, // 2020-02-02T12:34:56.123456Z
+        636467696123456L, // 1990-03-03T12:34:56.123456Z
+        165846896123456L // 1975-04-04T12:34:56.123456Z
+      };
+
+  public static final Long[] EXPECTED_TIMESTAMPS_MICROS =
+      new Long[] {
+        1735734896123456L, // 2025-01-01T12:34:56.123456Z
+        1580646896123456L, // 2020-02-02T12:34:56.123456Z
+        636467696123456L, // 1990-03-03T12:34:56.123456Z
+        165846896123456L // 1975-04-04T12:34:56.123456Z
+      };
+
+  public static ServiceAccountCredentials loadCredentials(String credentialFile) {
+    try (InputStream keyStream = new ByteArrayInputStream(credentialFile.getBytes())) {
+      return ServiceAccountCredentials.fromStream(keyStream);
+    } catch (IOException e) {
+      fail("Couldn't create fake JSON credentials.");
+    }
+    return null;
+  }
+
+  public static class AppendCompleteCallback implements ApiFutureCallback<AppendRowsResponse> {
+    private final Object lock = new Object();
+    private int batchCount = 0;
+
+    public void onSuccess(AppendRowsResponse response) {
+      synchronized (lock) {
+        if (response.hasError()) {
+          System.out.format("Error: %s\n", response.getError());
+        } else {
+          ++batchCount;
+          System.out.format("Wrote batch %d\n", batchCount);
+        }
+      }
+    }
+
+    public void onFailure(Throwable throwable) {
+      System.out.format("Error: %s\n", throwable.toString());
+    }
+  }
+
+  /**
+   * Reads all the rows from the specified table.
+   *
+   * <p>For every row, the consumer is called for processing.
+   *
+   * @param table
+   * @param snapshotInMillis Optional. If specified, all rows up to timestamp will be returned.
+   * @param filter Optional. If specified, it will be used to restrict returned data.
+   * @param consumer that receives all Avro rows.
+   * @throws IOException
+   */
+  public static void processRowsAtSnapshot(
+      BigQueryReadClient client,
+      String parentProjectId,
+      String table,
+      Long snapshotInMillis,
+      String filter,
+      SimpleRowReaderAvro.AvroRowConsumer consumer)
+      throws IOException {
+    Preconditions.checkNotNull(table);
+    Preconditions.checkNotNull(consumer);
+
+    CreateReadSessionRequest.Builder createSessionRequestBuilder =
+        CreateReadSessionRequest.newBuilder()
+            .setParent(parentProjectId)
+            .setMaxStreamCount(1)
+            .setReadSession(
+                ReadSession.newBuilder().setTable(table).setDataFormat(DataFormat.AVRO).build());
+
+    if (snapshotInMillis != null) {
+      createSessionRequestBuilder
+          .getReadSessionBuilder()
+          .setTableModifiers(
+              ReadSession.TableModifiers.newBuilder()
+                  .setSnapshotTime(Timestamps.fromMillis(snapshotInMillis))
+                  .build());
+    }
+
+    if (filter != null && !filter.isEmpty()) {
+      createSessionRequestBuilder
+          .getReadSessionBuilder()
+          .setReadOptions(
+              ReadSession.TableReadOptions.newBuilder().setRowRestriction(filter).build());
+    }
+
+    ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+    assertEquals(
+        String.format(
+            "Did not receive expected number of streams for table '%s' CreateReadSession"
+                + " response:%n%s",
+            table, session.toString()),
+        1,
+        session.getStreamsCount());
+
+    ReadRowsRequest readRowsRequest =
+        ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();
+
+    SimpleRowReaderAvro reader =
+        new SimpleRowReaderAvro(new Schema.Parser().parse(session.getAvroSchema().getSchema()));
+
+    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
+    for (ReadRowsResponse response : stream) {
+      reader.processRows(response.getAvroRows(), consumer);
+    }
+  }
+
+  /**
+   * Reads all the rows from the specified table and returns a list as generic Avro records.
+   *
+   * @param table
+   * @param filter Optional. If specified, it will be used to restrict returned data.
+   * @return
+   */
+  public static List<GenericData.Record> readAllRows(
+      BigQueryReadClient client, String parentProjectId, String table, String filter)
+      throws IOException {
+    final List<GenericData.Record> rows = new ArrayList<>();
+    processRowsAtSnapshot(
+        client,
+        parentProjectId,
+        /* table= */ table,
+        /* snapshotInMillis= */ null,
+        /* filter= */ filter,
+        (SimpleRowReaderAvro.AvroRowConsumer)
+            record -> {
+              // clone the record since that reference will be reused by the reader.
+              rows.add(new GenericRecordBuilder(record).build());
+            });
+    return rows;
+  }
+}

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/SimpleRowReaderArrow.java → google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/SimpleRowReaderArrow.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.google.cloud.bigquery.storage.v1.it;
+package com.google.cloud.bigquery.storage.v1.it.util;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -23,7 +23,6 @@
 import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
 import com.google.cloud.bigquery.storage.v1.ArrowSchema;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.time.LocalDateTime;
 import java.util.ArrayList;
@@ -50,17 +49,34 @@ public interface ArrowBatchConsumer {
     void accept(VectorSchemaRoot root);
   }
 
+  public static class ArrowTimestampBatchConsumer implements ArrowBatchConsumer {
+    private final List<Long> expectedTimestampValues;
+
+    public ArrowTimestampBatchConsumer(List<Long> expectedTimestampValues) {
+      this.expectedTimestampValues = expectedTimestampValues;
+    }
+
+    @Override
+    public void accept(VectorSchemaRoot root) {
+      FieldVector timestampFieldVector = root.getVector("timestamp");
+      int count = timestampFieldVector.getValueCount();
+      for (int i = 0; i < count; i++) {
+        long value = (Long) timestampFieldVector.getObject(i);
+        assertThat(value).isEqualTo(expectedTimestampValues.get(i));
+      }
+    }
+  }
+
   /** ArrowRangeBatchConsumer accepts batch Arrow data and validate the range values. */
   public static class ArrowRangeBatchConsumer implements ArrowBatchConsumer {
-
-    private final ImmutableMap<String, Range> expectedRangeDateValues;
-    private final ImmutableMap<String, Range> expectedRangeDatetimeValues;
-    private final ImmutableMap<String, Range> expectedRangeTimestampValues;
+    private final Map<String, Range> expectedRangeDateValues;
+    private final Map<String, Range> expectedRangeDatetimeValues;
+    private final Map<String, Range> expectedRangeTimestampValues;
 
     public ArrowRangeBatchConsumer(
-        ImmutableMap<String, Range> expectedRangeDateValues,
-        ImmutableMap<String, Range> expectedRangeDatetimeValues,
-        ImmutableMap<String, Range> expectedRangeTimestampValues) {
+        Map<String, Range> expectedRangeDateValues,
+        Map<String, Range> expectedRangeDatetimeValues,
+        Map<String, Range> expectedRangeTimestampValues) {
       this.expectedRangeDateValues = expectedRangeDateValues;
       this.expectedRangeDatetimeValues = expectedRangeDatetimeValues;
       this.expectedRangeTimestampValues = expectedRangeTimestampValues;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/SimpleRowReaderAvro.java → google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/SimpleRowReaderAvro.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.google.cloud.bigquery.storage.v1.it;
+package com.google.cloud.bigquery.storage.v1.it.util;
 
 import com.google.cloud.bigquery.storage.v1.AvroRows;
 import com.google.common.base.Preconditions;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/WriteRetryTestUtil.java → google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/WriteRetryTestUtil.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.google.cloud.bigquery.storage.v1.it;
+package com.google.cloud.bigquery.storage.v1.it.util;
 
 import static org.junit.Assert.assertFalse;
 

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

@@ -244,7 +244,7 @@ public void testSimpleRead() {
   @Test
   public void testSimpleReadArrow() {
     String table =
-        com.google.cloud.bigquery.storage.v1.it.BigQueryResource.formatTableResource(
+        com.google.cloud.bigquery.storage.v1.it.util.BigQueryResource.formatTableResource(
             /* projectId= */ "bigquery-public-data",
             /* datasetId= */ "samples",
             /* tableId= */ "shakespeare");
@@ -308,7 +308,7 @@ public void testRangeType() throws InterruptedException {
     bigquery.query(createTable);
 
     String table =
-        com.google.cloud.bigquery.storage.v1.it.BigQueryResource.formatTableResource(
+        com.google.cloud.bigquery.storage.v1.it.util.BigQueryResource.formatTableResource(
             /* projectId= */ ServiceOptions.getDefaultProjectId(),
             /* datasetId= */ DATASET,
             /* tableId= */ tableId.getTable());