Add 'Table.mutation_timeout', allowing override of config timeouts. (#7424)

Author: juan-rael

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

bigtable/google/cloud/bigtable/instance.py

@@ -538,7 +538,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:
@@ -556,7 +556,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,

bigtable/tests/unit/test_table.py

@@ -1205,59 +1205,6 @@ def test_callable_retry(self):
  )
  self.assertEqual(result, expected_result)
 
- def test_callable_retry_timeout(self):
- from google.cloud.bigtable.row import DirectRow
- from google.cloud.bigtable.table import DEFAULT_RETRY
- from google.cloud.bigtable_v2.gapic import bigtable_client
- from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client
-
- # Setup:
- # - Mutate 2 rows.
- # Action:
- # - Initial attempt will mutate all 2 rows.
- # Expectation:
- # - Both rows always return retryable errors.
- # - google.api_core.Retry should keep retrying.
- # - Check MutateRows is called multiple times.
- # - By the time deadline is reached, statuses should be
- # [retryable, retryable]
-
- data_api = bigtable_client.BigtableClient(mock.Mock())
- table_api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
- credentials = _make_credentials()
- client = self._make_client(
- project="project-id", credentials=credentials, admin=True
- )
- client._table_data_client = data_api
- client._table_admin_client = table_api
- instance = client.instance(instance_id=self.INSTANCE_ID)
- table = self._make_table(self.TABLE_ID, instance)
-
- row_1 = DirectRow(row_key=b"row_key", table=table)
- row_1.set_cell("cf", b"col", b"value1")
- row_2 = DirectRow(row_key=b"row_key_2", table=table)
- row_2.set_cell("cf", b"col", b"value2")
-
- response = self._make_responses([self.RETRYABLE_1, self.RETRYABLE_1])
-
- # Patch the stub used by the API method.
- inner_api_calls = client._table_data_client._inner_api_calls
- inner_api_calls["mutate_rows"] = mock.Mock(return_value=[response])
-
- retry = DEFAULT_RETRY.with_delay(
- initial=0.1, maximum=0.2, multiplier=2.0
- ).with_deadline(0.5)
- worker = self._make_worker(client, table.name, [row_1, row_2])
- statuses = worker(retry=retry)
-
- result = [status.code for status in statuses]
- expected_result = [self.RETRYABLE_1, self.RETRYABLE_1]
-
- self.assertTrue(
- client._table_data_client._inner_api_calls["mutate_rows"].call_count > 1
- )
- self.assertEqual(result, expected_result)
-
  def test_do_mutate_retryable_rows_empty_rows(self):
  from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client