use retry settings from config

Author: larkee

google/cloud/spanner_v1/database.py

@@ -23,7 +23,7 @@
 import warnings
 
 from google.api_core.client_options import ClientOptions
-from google.api_core.retry import Retry, if_exception_type
+from google.api_core.retry import if_exception_type
 import google.auth.credentials
 from google.protobuf.struct_pb2 import Struct
 from google.cloud.exceptions import NotFound
@@ -468,8 +468,9 @@ def execute_pdml():
 
  return result_set.stats.row_count_lower_bound
 
- retry = Retry(predicate=if_exception_type(Aborted))
- return retry(execute_pdml)()
+ retry_config = api._method_configs["ExecuteStreamingSql"].retry
+
+ return _retry_on_aborted(execute_pdml, retry_config)()
 
  def session(self, labels=None):
  """Factory to create a session for this database.
@@ -1029,3 +1030,19 @@ def __init__(self, source_type, backup_info):
  @classmethod
  def from_pb(cls, pb):
  return cls(pb.source_type, pb.backup_info)
+
+
+def _retry_on_aborted(func, retry_config):
+ """Helper for :meth:`Database.execute_partitioned_dml`.
+
+ Wrap function in a Retry that will retry on Aborted exceptions
+ with the retry config specified.
+
+ :type func: callable
+ :param func: the function to be retried on Aborted exceptions
+
+ :type retry_config: Retry
+ :param retry_config: retry object with the settings to be used
+ """
+ retry = retry_config.with_predicate(if_exception_type(Aborted))
+ return retry(func)

tests/unit/test_database.py

@@ -975,6 +975,7 @@ def _execute_partitioned_dml_helper(
  self, dml, params=None, param_types=None, query_options=None, retried=False
  ):
  from google.api_core.exceptions import Aborted
+ from google.api_core.retry import Retry
  from google.protobuf.struct_pb2 import Struct
  from google.cloud.spanner_v1.proto.result_set_pb2 import (
  PartialResultSet,
@@ -990,6 +991,10 @@ def _execute_partitioned_dml_helper(
  _merge_query_options,
  )
 
+ import collections
+
+ MethodConfig = collections.namedtuple("MethodConfig", ["retry"])
+
  transaction_pb = TransactionPB(id=self.TRANSACTION_ID)
 
  stats_pb = ResultSetStats(row_count_lower_bound=2)
@@ -1003,6 +1008,7 @@ def _execute_partitioned_dml_helper(
  pool.put(session)
  database = self._make_one(self.DATABASE_ID, instance, pool=pool)
  api = database._spanner_api = self._make_spanner_api()
+ api._method_configs = {"ExecuteStreamingSql": MethodConfig(retry=Retry())}
  if retried:
  retry_transaction_pb = TransactionPB(id=self.RETRY_TRANSACTION_ID)
  api.begin_transaction.side_effect = [transaction_pb, retry_transaction_pb]