Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
fix: set timeouts for BatchGetDocuments/RunQuery (#799)
  • Loading branch information
schmidt-sebastian committed Oct 27, 2021
1 parent ec07294 commit 6cd2a45
Show file tree
Hide file tree
Showing 5 changed files with 61 additions and 1 deletion.
Expand Up @@ -132,6 +132,9 @@ public GrpcFirestoreRpc(final FirestoreOptions options) throws IOException {
FirestoreStubSettings.Builder firestoreBuilder =
FirestoreStubSettings.newBuilder(clientContext)
.applyToAllUnaryMethods(retrySettingsSetter);
// Manually apply the retry settings to streaming methods
firestoreBuilder.runQuerySettings().setRetrySettings(options.getRetrySettings());
firestoreBuilder.batchGetDocumentsSettings().setRetrySettings(options.getRetrySettings());
firestoreStub = GrpcFirestoreStub.create(firestoreBuilder.build());
} catch (Exception e) {
throw new IOException(e);
Expand Down
Expand Up @@ -14,7 +14,7 @@
* limitations under the License.
*/

package com.google.cloud.firestore;
package com.google.cloud.firestore.it;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
Expand All @@ -24,6 +24,14 @@

import com.google.api.core.ApiFuture;
import com.google.cloud.Timestamp;
import com.google.cloud.firestore.BulkWriter;
import com.google.cloud.firestore.CollectionReference;
import com.google.cloud.firestore.DocumentReference;
import com.google.cloud.firestore.DocumentSnapshot;
import com.google.cloud.firestore.Firestore;
import com.google.cloud.firestore.FirestoreOptions;
import com.google.cloud.firestore.LocalFirestoreHelper;
import com.google.cloud.firestore.WriteResult;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.Collections;
Expand Down
Expand Up @@ -58,7 +58,10 @@
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

@RunWith(JUnit4.class)
public final class ITQueryWatchTest {

private static Firestore firestore;
Expand Down
Expand Up @@ -32,7 +32,10 @@
import org.junit.Test;
import org.junit.rules.TestName;
import org.junit.rules.Timeout;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

@RunWith(JUnit4.class)
public class ITShutdownTest {
@Rule public final Timeout timeout = new Timeout(5, TimeUnit.SECONDS);
@Rule public TestName testName = new TestName();
Expand Down
Expand Up @@ -33,12 +33,14 @@
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.core.SettableApiFuture;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.ApiStreamObserver;
import com.google.cloud.Timestamp;
import com.google.cloud.firestore.BulkWriter;
Expand Down Expand Up @@ -100,7 +102,11 @@
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.threeten.bp.Duration;

@RunWith(JUnit4.class)
public class ITSystemTest {

private static final double DOUBLE_EPSILON = 0.000001;
Expand Down Expand Up @@ -1761,4 +1767,41 @@ public void testRecursiveDeleteWithCustomBulkWriterInstance() throws Exception {
assertEquals(0, countCollectionChildren(randomColl));
assertEquals(6, callbackCount[0]);
}

@Test
public void testEnforcesTimeouts() throws Exception {
FirestoreOptions firestoreOptions =
FirestoreOptions.newBuilder()
.setRetrySettings(
RetrySettings.newBuilder()
.setMaxRpcTimeout(Duration.ofMillis(1))
.setTotalTimeout(Duration.ofMillis(1))
.setInitialRpcTimeout(Duration.ofMillis(1))
.build())
.build();
firestore = firestoreOptions.getService();
CollectionReference collection = firestore.collection("timeout");

// RunQuery
assertThrows(ExecutionException.class, () -> collection.get().get());
// CommitRequest
assertThrows(ExecutionException.class, () -> collection.add(map()).get());
// BulkCommit
assertThrows(
ExecutionException.class,
() -> {
BulkWriter bulkWriter = firestore.bulkWriter();
ApiFuture<WriteResult> op = bulkWriter.set(collection.document(), map());
bulkWriter.close();
op.get();
});
// BatchGetDocuments
assertThrows(ExecutionException.class, () -> collection.document().get().get());
// ListDocuments
assertThrows(FirestoreException.class, () -> collection.listDocuments().iterator().hasNext());
// ListCollections
assertThrows(
FirestoreException.class,
() -> collection.document().listCollections().iterator().hasNext());
}
}

0 comments on commit 6cd2a45

Please sign in to comment.