Custom Timeout in Table Mutation

Author: Juan Rael

bigtable/docs/snippets_table.py

@@ -39,22 +39,24 @@
 from google.cloud.bigtable import column_family
 
 
-INSTANCE_ID = "snippet-" + unique_resource_id('-')
-CLUSTER_ID = "clus-1-" + unique_resource_id('-')
-TABLE_ID = "tabl-1-" + unique_resource_id('-')
-COLUMN_FAMILY_ID = "col_fam_id-" + unique_resource_id('-')
-LOCATION_ID = 'us-central1-f'
-ALT_LOCATION_ID = 'us-central1-a'
+INSTANCE_ID = "snippet-" + unique_resource_id("-")
+CLUSTER_ID = "clus-1-" + unique_resource_id("-")
+TABLE_ID = "tabl-1-" + unique_resource_id("-")
+COLUMN_FAMILY_ID = "col_fam_id-" + unique_resource_id("-")
+LOCATION_ID = "us-central1-f"
+ALT_LOCATION_ID = "us-central1-a"
 PRODUCTION = enums.Instance.Type.PRODUCTION
 SERVER_NODES = 3
 STORAGE_TYPE = enums.StorageType.SSD
-LABEL_KEY = u'python-snippet'
-LABEL_STAMP = datetime.datetime.utcnow() \
- .replace(microsecond=0, tzinfo=UTC,) \
- .strftime("%Y-%m-%dt%H-%M-%S")
+LABEL_KEY = u"python-snippet"
+LABEL_STAMP = (
+ datetime.datetime.utcnow()
+ .replace(microsecond=0, tzinfo=UTC)
+ .strftime("%Y-%m-%dt%H-%M-%S")
+)
 LABELS = {LABEL_KEY: str(LABEL_STAMP)}
-COL_NAME1 = b'col-name1'
-CELL_VAL1 = b'cell-val'
+COL_NAME1 = b"col-name1"
+CELL_VAL1 = b"cell-val"
 
 
 class Config(object):
@@ -63,28 +65,30 @@ class Config(object):
  This is a mutable stand-in to allow test set-up to modify
  global state.
  """
+
  CLIENT = None
  INSTANCE = None
  TABLE = None
 
 
 def setup_module():
  client = Config.CLIENT = Client(admin=True)
- Config.INSTANCE = client.instance(INSTANCE_ID,
- instance_type=PRODUCTION,
- labels=LABELS)
- cluster = Config.INSTANCE.cluster(CLUSTER_ID,
- location_id=LOCATION_ID,
- serve_nodes=SERVER_NODES,
- default_storage_type=STORAGE_TYPE)
+ Config.INSTANCE = client.instance(
+ INSTANCE_ID, instance_type=PRODUCTION, labels=LABELS
+ )
+ cluster = Config.INSTANCE.cluster(
+ CLUSTER_ID,
+ location_id=LOCATION_ID,
+ serve_nodes=SERVER_NODES,
+ default_storage_type=STORAGE_TYPE,
+ )
  operation = Config.INSTANCE.create(clusters=[cluster])
  # We want to make sure the operation completes.
  operation.result(timeout=100)
  Config.TABLE = Config.INSTANCE.table(TABLE_ID)
  Config.TABLE.create()
  gc_rule = column_family.MaxVersionsGCRule(2)
- column_family1 = Config.TABLE.column_family(COLUMN_FAMILY_ID,
- gc_rule=gc_rule)
+ column_family1 = Config.TABLE.column_family(COLUMN_FAMILY_ID, gc_rule=gc_rule)
  column_family1.create()
 
 
@@ -108,7 +112,7 @@ def test_bigtable_create_table():
  table2 = instance.table("table_id2")
  # Define the GC policy to retain only the most recent 2 versions.
  max_versions_rule = column_family.MaxVersionsGCRule(2)
- table2.create(column_families={'cf1': max_versions_rule})
+ table2.create(column_families={"cf1": max_versions_rule})
 
  # [END bigtable_create_table]
  assert table1.exists()
@@ -126,14 +130,13 @@ def test_bigtable_sample_row_keys():
 
  table = instance.table("table_id1_samplerow")
  # [END bigtable_sample_row_keys]
- initial_split_keys = [b'split_key_1', b'split_key_10',
- b'split_key_20']
+ initial_split_keys = [b"split_key_1", b"split_key_10", b"split_key_20"]
  table.create(initial_split_keys=initial_split_keys)
  # [START bigtable_sample_row_keys]
  data = table.sample_row_keys()
  actual_keys, offset = zip(*[(rk.row_key, rk.offset_bytes) for rk in data])
  # [END bigtable_sample_row_keys]
- initial_split_keys.append(b'')
+ initial_split_keys.append(b"")
  assert list(actual_keys) == initial_split_keys
  table.delete()
 
@@ -145,23 +148,29 @@ def test_bigtable_write_read_drop_truncate():
  client = Client(admin=True)
  instance = client.instance(INSTANCE_ID)
  table = instance.table(TABLE_ID)
- row_keys = [b'row_key_1', b'row_key_2', b'row_key_3', b'row_key_4',
- b'row_key_20', b'row_key_22', b'row_key_200']
- col_name = b'col-name1'
+ row_keys = [
+ b"row_key_1",
+ b"row_key_2",
+ b"row_key_3",
+ b"row_key_4",
+ b"row_key_20",
+ b"row_key_22",
+ b"row_key_200",
+ ]
+ col_name = b"col-name1"
  rows = []
  for i, row_key in enumerate(row_keys):
- value = 'value_{}'.format(i).encode()
+ value = "value_{}".format(i).encode()
  row = table.row(row_key)
- row.set_cell(COLUMN_FAMILY_ID,
- col_name,
- value,
- timestamp=datetime.datetime.utcnow())
+ row.set_cell(
+ COLUMN_FAMILY_ID, col_name, value, timestamp=datetime.datetime.utcnow()
+ )
  rows.append(row)
  response = table.mutate_rows(rows)
  # validate that all rows written successfully
  for i, status in enumerate(response):
  if status.code is not 0:
- print('Row number {} failed to write'.format(i))
+ print("Row number {} failed to write".format(i))
  # [END bigtable_mutate_rows]
  assert len(response) == len(rows)
  # [START bigtable_read_row]
@@ -170,10 +179,10 @@ def test_bigtable_write_read_drop_truncate():
  client = Client(admin=True)
  instance = client.instance(INSTANCE_ID)
  table = instance.table(TABLE_ID)
- row_key = 'row_key_1'
+ row_key = "row_key_1"
  row = table.read_row(row_key)
  # [END bigtable_read_row]
- assert row.row_key.decode('utf-8') == row_key
+ assert row.row_key.decode("utf-8") == row_key
  # [START bigtable_read_rows]
  from google.cloud.bigtable import Client
 
@@ -192,13 +201,12 @@ def test_bigtable_write_read_drop_truncate():
  client = Client(admin=True)
  instance = client.instance(INSTANCE_ID)
  table = instance.table(TABLE_ID)
- row_key_prefix = b'row_key_2'
+ row_key_prefix = b"row_key_2"
  table.drop_by_prefix(row_key_prefix, timeout=200)
  # [END bigtable_drop_by_prefix]
- dropped_row_keys = [b'row_key_2', b'row_key_20',
- b'row_key_22', b'row_key_200']
+ dropped_row_keys = [b"row_key_2", b"row_key_20", b"row_key_22", b"row_key_200"]
  for row in table.read_rows():
- assert row.row_key.decode('utf-8') not in dropped_row_keys
+ assert row.row_key.decode("utf-8") not in dropped_row_keys
 
  # [START bigtable_truncate_table]
  from google.cloud.bigtable import Client
@@ -226,26 +234,31 @@ def test_bigtable_mutations_batcher():
 
  # Below code will be used while creating batcher.py snippets.
  # So not removing this code as of now.
- row_keys = [b'row_key_1', b'row_key_2', b'row_key_3', b'row_key_4',
- b'row_key_20', b'row_key_22', b'row_key_200']
- column_name = 'column_name'.encode()
+ row_keys = [
+ b"row_key_1",
+ b"row_key_2",
+ b"row_key_3",
+ b"row_key_4",
+ b"row_key_20",
+ b"row_key_22",
+ b"row_key_200",
+ ]
+ column_name = "column_name".encode()
  # Add a single row
  row_key = row_keys[0]
  row = table.row(row_key)
- row.set_cell(COLUMN_FAMILY_ID,
- column_name,
- 'value-0',
- timestamp=datetime.datetime.utcnow())
+ row.set_cell(
+ COLUMN_FAMILY_ID, column_name, "value-0", timestamp=datetime.datetime.utcnow()
+ )
  batcher.mutate(row)
  # Add a collections of rows
  rows = []
  for i in range(1, len(row_keys)):
  row = table.row(row_keys[i])
- value = 'value_{}'.format(i).encode()
- row.set_cell(COLUMN_FAMILY_ID,
- column_name,
- value,
- timestamp=datetime.datetime.utcnow())
+ value = "value_{}".format(i).encode()
+ row.set_cell(
+ COLUMN_FAMILY_ID, column_name, value, timestamp=datetime.datetime.utcnow()
+ )
  rows.append(row)
  batcher.mutate_rows(rows)
  # batcher will flush current batch if it
@@ -287,6 +300,7 @@ def test_bigtable_list_tables():
 
 def test_bigtable_table_name():
  import re
+
  # [START bigtable_table_name]
  from google.cloud.bigtable import Client
 
@@ -296,9 +310,11 @@ def test_bigtable_table_name():
  table = instance.table(TABLE_ID)
  table_name = table.name
  # [END bigtable_table_name]
- _table_name_re = re.compile(r'^projects/(?P<project>[^/]+)/'
- r'instances/(?P<instance>[^/]+)/tables/'
- r'(?P<table_id>[_a-zA-Z0-9][-_.a-zA-Z0-9]*)$')
+ _table_name_re = re.compile(
+ r"^projects/(?P<project>[^/]+)/"
+ r"instances/(?P<instance>[^/]+)/tables/"
+ r"(?P<table_id>[_a-zA-Z0-9][-_.a-zA-Z0-9]*)$"
+ )
  assert _table_name_re.match(table_name)
 
 
@@ -368,7 +384,7 @@ def test_bigtable_table_row():
  instance = client.instance(INSTANCE_ID)
  table = instance.table(TABLE_ID)
 
- row_keys = [b'row_key_1', b'row_key_2']
+ row_keys = [b"row_key_1", b"row_key_2"]
  row1_obj = table.row(row_keys[0])
  row2_obj = table.row(row_keys[1])
  # [END bigtable_table_row]
@@ -387,5 +403,5 @@ def test_bigtable_table_row():
  table.truncate(timeout=300)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
  pytest.main()

bigtable/google/cloud/bigtable/batcher.py

@@ -56,7 +56,9 @@ class MutationsBatcher(object):
  (5 MB).
  """
 
- def __init__(self, table, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES):
+ def __init__(
+ self, table, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES
+ ):
  self.rows = []
  self.total_mutation_count = 0
  self.total_size = 0
@@ -144,7 +146,7 @@ def mutate_rows(self, rows):
 
  def flush(self):
  """ Sends the current. batch to Cloud Bigtable. """
- if len(self.rows) is not 0:
+ if len(self.rows) != 0:
  self.table.mutate_rows(self.rows)
  self.total_mutation_count = 0
  self.total_size = 0

bigtable/google/cloud/bigtable/instance.py

@@ -553,7 +553,7 @@ def list_clusters(self):
  clusters = [Cluster.from_pb(cluster, self) for cluster in resp.clusters]
  return clusters, resp.failed_locations
 
- def table(self, table_id, app_profile_id=None):
+ def table(self, table_id, mutation_timeout=None, app_profile_id=None):
  """Factory to create a table associated with this instance.
 
  For example:
@@ -571,7 +571,12 @@ def table(self, table_id, app_profile_id=None):
  :rtype: :class:`Table <google.cloud.bigtable.table.Table>`
  :returns: The table owned by this instance.
  """
- return Table(table_id, self, app_profile_id=app_profile_id)
+ return Table(
+ table_id,
+ self,
+ app_profile_id=app_profile_id,
+ mutation_timeout=mutation_timeout,
+ )
 
  def list_tables(self):
  """List the tables in this instance.

bigtable/google/cloud/bigtable/table.py

@@ -17,6 +17,7 @@
 
 from grpc import StatusCode
 
+from google.api_core import timeout
 from google.api_core.exceptions import RetryError
 from google.api_core.exceptions import NotFound
 from google.api_core.retry import if_exception_type
@@ -100,10 +101,11 @@ class Table(object):
  :param app_profile_id: (Optional) The unique name of the AppProfile.
  """
 
- def __init__(self, table_id, instance, app_profile_id=None):
+ def __init__(self, table_id, instance, mutation_timeout=None, app_profile_id=None):
  self.table_id = table_id
  self._instance = instance
  self._app_profile_id = app_profile_id
+ self.mutation_timeout = mutation_timeout
 
  @property
  def name(self):
@@ -503,7 +505,11 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY):
  sent. These will be in the same order as the `rows`.
  """
  retryable_mutate_rows = _RetryableMutateRowsWorker(
- self._instance._client, self.name, rows, app_profile_id=self._app_profile_id
+ self._instance._client,
+ self.name,
+ rows,
+ app_profile_id=self._app_profile_id,
+ timeout=self.mutation_timeout,
  )
  return retryable_mutate_rows(retry=retry)
 
@@ -658,12 +664,13 @@ class _RetryableMutateRowsWorker(object):
  )
  # pylint: enable=unsubscriptable-object
 
- def __init__(self, client, table_name, rows, app_profile_id=None):
+ def __init__(self, client, table_name, rows, app_profile_id=None, timeout=None):
  self.client = client
  self.table_name = table_name
  self.rows = rows
  self.app_profile_id = app_profile_id
  self.responses_statuses = [None] * len(self.rows)
+ self.timeout = timeout
 
  def __call__(self, retry=DEFAULT_RETRY):
  """Attempt to mutate all rows and retry rows with transient errors.
@@ -729,7 +736,10 @@ def _do_mutate_retryable_rows(self):
  inner_api_calls = data_client._inner_api_calls
  if "mutate_rows" not in inner_api_calls:
  default_retry = (data_client._method_configs["MutateRows"].retry,)
- default_timeout = data_client._method_configs["MutateRows"].timeout
+ if self.timeout is None:
+ default_timeout = data_client._method_configs["MutateRows"].timeout
+ else:
+ default_timeout = timeout.ExponentialTimeout(deadline=self.timeout)
  data_client._inner_api_calls["mutate_rows"] = wrap_method(
  data_client.transport.mutate_rows,
  default_retry=default_retry,