[bigtable] fix: wrap sample invocations with retries (GoogleCloudPlatform#3494)

Takashi Matsuo · web-flow · commit bb83cea15029 · 2020-04-24T11:59:21.000-07:00
fix GoogleCloudPlatform#3070 Also added `BIGTABLE_INSTANCE` to testing/test-env.tmpl.sh
diff --git a/bigtable/snippets/writes/requirements-test.txt b/bigtable/snippets/writes/requirements-test.txt
@@ -1 +1,2 @@
+backoff==1.10.0
 pytest==5.3.2
diff --git a/bigtable/snippets/writes/writes_test.py b/bigtable/snippets/writes/writes_test.py
@@ -16,6 +16,8 @@
 import uuid
 import pytest
 
+import backoff
+from google.api_core.exceptions import DeadlineExceeded
 from google.cloud import bigtable
 
 from .write_batch import write_batch
@@ -55,22 +57,37 @@ def table_id(bigtable_instance):
 
 
 def test_writes(capsys, table_id):
-    write_simple(PROJECT, BIGTABLE_INSTANCE, table_id)
 
+    # `row.commit()` sometimes ends up with DeadlineExceeded, so now
+    # we put retries with a hard deadline.
+    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    def _write_simple():
+        write_simple(PROJECT, BIGTABLE_INSTANCE, table_id)
+
+    _write_simple()
     out, _ = capsys.readouterr()
     assert 'Successfully wrote row' in out
 
-    write_increment(PROJECT, BIGTABLE_INSTANCE, table_id)
+    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    def _write_increment():
+        write_increment(PROJECT, BIGTABLE_INSTANCE, table_id)
 
+    _write_increment()
     out, _ = capsys.readouterr()
     assert 'Successfully updated row' in out
 
-    write_conditional(PROJECT, BIGTABLE_INSTANCE, table_id)
+    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    def _write_conditional():
+        write_conditional(PROJECT, BIGTABLE_INSTANCE, table_id)
 
+    _write_conditional()
     out, _ = capsys.readouterr()
     assert 'Successfully updated row\'s os_name' in out
 
-    write_batch(PROJECT, BIGTABLE_INSTANCE, table_id)
+    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    def _write_batch():
+        write_batch(PROJECT, BIGTABLE_INSTANCE, table_id)
 
+    _write_batch()
     out, _ = capsys.readouterr()
     assert 'Successfully wrote 2 rows' in out
diff --git a/testing/test-env.tmpl.sh b/testing/test-env.tmpl.sh
@@ -8,6 +8,7 @@ export CLOUD_STORAGE_BUCKET=$GCLOUD_PROJECT
 export API_KEY=
 export BIGTABLE_CLUSTER=bigtable-test
 export BIGTABLE_ZONE=us-central1-c
+export BIGTABLE_INSTANCE=
 export SPANNER_INSTANCE=
 export COMPOSER_LOCATION=us-central1
 export COMPOSER_ENVIRONMENT=
