feat: add retry and timeout for batch dml (#1107)

* feat(spanner): add retry, timeout for batch update * feat(spanner): add samples for retry, timeout * feat(spanner): update unittest * feat(spanner): update comments * feat(spanner): update code for retry * feat(spanner): update comment

Author: harshachinta

google/cloud/spanner_v1/transaction.py

@@ -410,7 +410,14 @@ def execute_update(
 
  return response.stats.row_count_exact
 
- def batch_update(self, statements, request_options=None):
+ def batch_update(
+ self,
+ statements,
+ request_options=None,
+ *,
+ retry=gapic_v1.method.DEFAULT,
+ timeout=gapic_v1.method.DEFAULT,
+ ):
  """Perform a batch of DML statements via an ``ExecuteBatchDml`` request.
 
  :type statements:
@@ -431,6 +438,12 @@ def batch_update(self, statements, request_options=None):
  If a dict is provided, it must be of the same form as the protobuf
  message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
 
+ :type retry: :class:`~google.api_core.retry.Retry`
+ :param retry: (Optional) The retry settings for this request.
+
+ :type timeout: float
+ :param timeout: (Optional) The timeout for this request.
+
  :rtype:
  Tuple(status, Sequence[int])
  :returns:
@@ -486,6 +499,8 @@ def batch_update(self, statements, request_options=None):
  api.execute_batch_dml,
  request=request,
  metadata=metadata,
+ retry=retry,
+ timeout=timeout,
  )
 
  if self._transaction_id is None:

samples/samples/snippets.py

@@ -3017,6 +3017,51 @@ def directed_read_options(
  # [END spanner_directed_read]
 
 
+def set_custom_timeout_and_retry(instance_id, database_id):
+ """Executes a snapshot read with custom timeout and retry."""
+ # [START spanner_set_custom_timeout_and_retry]
+ from google.api_core import retry
+ from google.api_core import exceptions as core_exceptions
+
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ retry = retry.Retry(
+ # Customize retry with an initial wait time of 500 milliseconds.
+ initial=0.5,
+ # Customize retry with a maximum wait time of 16 seconds.
+ maximum=16,
+ # Customize retry with a wait time multiplier per iteration of 1.5.
+ multiplier=1.5,
+ # Customize retry with a timeout on
+ # how long a certain RPC may be retried in
+ # case the server returns an error.
+ timeout=60,
+ # Configure which errors should be retried.
+ predicate=retry.if_exception_type(
+ core_exceptions.ServiceUnavailable,
+ ),
+ )
+
+ # Set a custom retry and timeout setting.
+ with database.snapshot() as snapshot:
+ results = snapshot.execute_sql(
+ "SELECT SingerId, AlbumId, AlbumTitle FROM Albums",
+ # Set custom retry setting for this request
+ retry=retry,
+ # Set custom timeout of 60 seconds for this request
+ timeout=60,
+ )
+
+ for row in results:
+ print("SingerId: {}, AlbumId: {}, AlbumTitle: {}".format(*row))
+
+ # [END spanner_set_custom_timeout_and_retry]
+
+
 if __name__ == "__main__": # noqa: C901
  parser = argparse.ArgumentParser(
  description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
@@ -3157,6 +3202,9 @@ def directed_read_options(
  )
  enable_fine_grained_access_parser.add_argument("--title", default="condition title")
  subparsers.add_parser("directed_read_options", help=directed_read_options.__doc__)
+ subparsers.add_parser(
+ "set_custom_timeout_and_retry", help=set_custom_timeout_and_retry.__doc__
+ )
 
  args = parser.parse_args()
 
@@ -3290,3 +3338,5 @@ def directed_read_options(
  )
  elif args.command == "directed_read_options":
  directed_read_options(args.instance_id, args.database_id)
+ elif args.command == "set_custom_timeout_and_retry":
+ set_custom_timeout_and_retry(args.instance_id, args.database_id)

samples/samples/snippets_test.py

@@ -859,3 +859,10 @@ def test_directed_read_options(capsys, instance_id, sample_database):
  snippets.directed_read_options(instance_id, sample_database.database_id)
  out, _ = capsys.readouterr()
  assert "SingerId: 1, AlbumId: 1, AlbumTitle: Total Junk" in out
+
+
+@pytest.mark.dependency(depends=["insert_data"])
+def test_set_custom_timeout_and_retry(capsys, instance_id, sample_database):
+ snippets.set_custom_timeout_and_retry(instance_id, sample_database.database_id)
+ out, _ = capsys.readouterr()
+ assert "SingerId: 1, AlbumId: 1, AlbumTitle: Total Junk" in out

tests/unit/test_spanner.py

@@ -556,6 +556,8 @@ def test_transaction_should_include_begin_with_first_batch_update(self):
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
 
  def test_transaction_should_use_transaction_id_if_error_with_first_batch_update(
@@ -574,6 +576,8 @@ def test_transaction_should_use_transaction_id_if_error_with_first_batch_update(
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
  self._execute_update_helper(transaction=transaction, api=api)
  api.execute_sql.assert_called_once_with(
@@ -715,6 +719,8 @@ def test_transaction_should_use_transaction_id_returned_by_first_read(self):
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
 
  def test_transaction_should_use_transaction_id_returned_by_first_batch_update(self):
@@ -729,6 +735,8 @@ def test_transaction_should_use_transaction_id_returned_by_first_batch_update(se
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
  self._read_helper(transaction=transaction, api=api)
  api.streaming_read.assert_called_once_with(
@@ -797,6 +805,8 @@ def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
 
  self.assertEqual(api.execute_sql.call_count, 2)
@@ -846,6 +856,8 @@ def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
 
  api.execute_batch_dml.assert_any_call(
@@ -854,6 +866,8 @@ def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=RETRY,
+ timeout=TIMEOUT,
  )
 
  self.assertEqual(api.execute_sql.call_count, 1)

tests/unit/test_transaction.py

@@ -662,7 +662,14 @@ def test_batch_update_other_error(self):
  with self.assertRaises(RuntimeError):
  transaction.batch_update(statements=[DML_QUERY])
 
- def _batch_update_helper(self, error_after=None, count=0, request_options=None):
+ def _batch_update_helper(
+ self,
+ error_after=None,
+ count=0,
+ request_options=None,
+ retry=gapic_v1.method.DEFAULT,
+ timeout=gapic_v1.method.DEFAULT,
+ ):
  from google.rpc.status_pb2 import Status
  from google.protobuf.struct_pb2 import Struct
  from google.cloud.spanner_v1 import param_types
@@ -716,7 +723,10 @@ def _batch_update_helper(self, error_after=None, count=0, request_options=None):
  request_options = RequestOptions(request_options)
 
  status, row_counts = transaction.batch_update(
- dml_statements, request_options=request_options
+ dml_statements,
+ request_options=request_options,
+ retry=retry,
+ timeout=timeout,
  )
 
  self.assertEqual(status, expected_status)
@@ -753,6 +763,8 @@ def _batch_update_helper(self, error_after=None, count=0, request_options=None):
  ("google-cloud-resource-prefix", database.name),
  ("x-goog-spanner-route-to-leader", "true"),
  ],
+ retry=retry,
+ timeout=timeout,
  )
 
  self.assertEqual(transaction._execute_sql_count, count + 1)
@@ -826,6 +838,15 @@ def test_batch_update_error(self):
 
  self.assertEqual(transaction._execute_sql_count, 1)
 
+ def test_batch_update_w_timeout_param(self):
+ self._batch_update_helper(timeout=2.0)
+
+ def test_batch_update_w_retry_param(self):
+ self._batch_update_helper(retry=gapic_v1.method.DEFAULT)
+
+ def test_batch_update_w_timeout_and_retry_params(self):
+ self._batch_update_helper(retry=gapic_v1.method.DEFAULT, timeout=2.0)
+
  def test_context_mgr_success(self):
  import datetime
  from google.cloud.spanner_v1 import CommitResponse