diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml
index 5d2607f2ac0e..cb4f796140c6 100644
--- a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml
+++ b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml
@@ -197,6 +197,11 @@
opentelemetry-sdk-trace
test
+
+ org.awaitility
+ awaitility
+ test
+
junit
junit
diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java
index 04191a38a45e..4e631baebe99 100644
--- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java
+++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java
@@ -18,6 +18,7 @@
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
+import static org.awaitility.Awaitility.await;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -26,6 +27,7 @@
import static org.junit.jupiter.api.Assertions.fail;
import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.rpc.NotFoundException;
import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.UnauthenticatedException;
import com.google.auth.oauth2.ServiceAccountCredentials;
@@ -81,6 +83,7 @@
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Logger;
import org.apache.avro.Conversions;
import org.apache.avro.LogicalTypes;
@@ -1230,7 +1233,19 @@ private void ProcessRowsAtSnapshot(
TableReadOptions.newBuilder().setRowRestriction(filter).build());
}
- ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+ CreateReadSessionRequest request = createSessionRequestBuilder.build();
+ AtomicReference sessionRef = new AtomicReference<>();
+ await()
+ .atMost(Duration.ofSeconds(10))
+ .pollInterval(Duration.ofSeconds(1))
+ // retry if the newly-created table has not yet fully propagated
+ .ignoreException(NotFoundException.class)
+ .until(
+ () -> {
+ sessionRef.set(client.createReadSession(request));
+ return true;
+ });
+ ReadSession session = sessionRef.get();
assertEquals(
1,
session.getStreamsCount(),
diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java
index f48d41acbd70..88b2f48b8b7f 100644
--- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java
+++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java
@@ -18,6 +18,7 @@
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
+import static org.awaitility.Awaitility.await;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -26,6 +27,7 @@
import static org.junit.jupiter.api.Assertions.fail;
import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.rpc.NotFoundException;
import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.UnauthenticatedException;
import com.google.auth.oauth2.ServiceAccountCredentials;
@@ -77,6 +79,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Logger;
import org.apache.avro.Conversions;
import org.apache.avro.LogicalTypes;
@@ -1211,7 +1214,19 @@ private void ProcessRowsAtSnapshot(
.setReadOptions(TableReadOptions.newBuilder().setRowRestriction(filter).build());
}
- ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+ CreateReadSessionRequest request = createSessionRequestBuilder.build();
+ AtomicReference sessionRef = new AtomicReference<>();
+ await()
+ .atMost(Duration.ofSeconds(10))
+ .pollInterval(Duration.ofSeconds(1))
+ // retry if the newly-created table has not yet fully propagated
+ .ignoreException(NotFoundException.class)
+ .until(
+ () -> {
+ sessionRef.set(client.createReadSession(request));
+ return true;
+ });
+ ReadSession session = sessionRef.get();
assertEquals(
1,
session.getStreamsCount(),