Retry streaming exceptions (this time for sure, Rocky!) (#4016)

* Add '_restart_on_unavailable' iterator wrapper.

  Tracks the 'resume_token', and issues restart after a 503.

* Strip knowledge of 'resume_token' from 'StreamedResultSet'.

* Remove 'resume_token' args from 'Snapshot' and 'Session' API surface:

  Retry handling will be done behind the scenes.

* Use '_restart_on_unavailable' wrapper in 'SRS.{read,execute_sql}.

Closes #3775.

spanner/google/cloud/spanner/session.py

@@ -165,8 +165,7 @@ def snapshot(self, **kw):
 
  return Snapshot(self, **kw)
 
- def read(self, table, columns, keyset, index='', limit=0,
- resume_token=b''):
+ def read(self, table, columns, keyset, index='', limit=0):
  """Perform a ``StreamingRead`` API request for rows in a table.
 
  :type table: str
@@ -185,17 +184,12 @@ def read(self, table, columns, keyset, index='', limit=0,
  :type limit: int
  :param limit: (Optional) maxiumn number of rows to return
 
- :type resume_token: bytes
- :param resume_token: token for resuming previously-interrupted read
-
  :rtype: :class:`~google.cloud.spanner.streamed.StreamedResultSet`
  :returns: a result set instance which can be used to consume rows.
  """
- return self.snapshot().read(
- table, columns, keyset, index, limit, resume_token)
+ return self.snapshot().read(table, columns, keyset, index, limit)
 
- def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
- resume_token=b''):
+ def execute_sql(self, sql, params=None, param_types=None, query_mode=None):
  """Perform an ``ExecuteStreamingSql`` API request.
 
  :type sql: str
@@ -216,14 +210,11 @@ def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
  :param query_mode: Mode governing return of results / query plan. See
  https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode1
 
- :type resume_token: bytes
- :param resume_token: token for resuming previously-interrupted query
-
  :rtype: :class:`~google.cloud.spanner.streamed.StreamedResultSet`
  :returns: a result set instance which can be used to consume rows.
  """
  return self.snapshot().execute_sql(
- sql, params, param_types, query_mode, resume_token)
+ sql, params, param_types, query_mode)
 
  def batch(self):
  """Factory to create a batch for this session.

spanner/google/cloud/spanner/snapshot.py

@@ -14,10 +14,13 @@
 
 """Model a set of read-only queries to a database as a snapshot."""
 
+import functools
+
 from google.protobuf.struct_pb2 import Struct
 from google.cloud.proto.spanner.v1.transaction_pb2 import TransactionOptions
 from google.cloud.proto.spanner.v1.transaction_pb2 import TransactionSelector
 
+from google.api.core.exceptions import ServiceUnavailable
 from google.cloud._helpers import _datetime_to_pb_timestamp
 from google.cloud._helpers import _timedelta_to_duration_pb
 from google.cloud.spanner._helpers import _make_value_pb
@@ -26,6 +29,36 @@
 from google.cloud.spanner.streamed import StreamedResultSet
 
 
+def _restart_on_unavailable(restart):
+ """Restart iteration after :exc:`.ServiceUnavailable`.
+
+ :type restart: callable
+ :param restart: curried function returning iterator
+ """
+ resume_token = ''
+ item_buffer = []
+ iterator = restart()
+ while True:
+ try:
+ for item in iterator:
+ item_buffer.append(item)
+ if item.resume_token:
+ resume_token = item.resume_token
+ break
+ except ServiceUnavailable:
+ del item_buffer[:]
+ iterator = restart(resume_token=resume_token)
+ continue
+
+ if len(item_buffer) == 0:
+ break
+
+ for item in item_buffer:
+ yield item
+
+ del item_buffer[:]
+
+
 class _SnapshotBase(_SessionWrapper):
  """Base class for Snapshot.
 
@@ -49,8 +82,7 @@ def _make_txn_selector(self): # pylint: disable=redundant-returns-doc
  """
  raise NotImplementedError
 
- def read(self, table, columns, keyset, index='', limit=0,
- resume_token=b''):
+ def read(self, table, columns, keyset, index='', limit=0):
  """Perform a ``StreamingRead`` API request for rows in a table.
 
  :type table: str
@@ -69,9 +101,6 @@ def read(self, table, columns, keyset, index='', limit=0,
  :type limit: int
  :param limit: (Optional) maxiumn number of rows to return
 
- :type resume_token: bytes
- :param resume_token: token for resuming previously-interrupted read
-
  :rtype: :class:`~google.cloud.spanner.streamed.StreamedResultSet`
  :returns: a result set instance which can be used to consume rows.
  :raises ValueError:
@@ -89,10 +118,13 @@ def read(self, table, columns, keyset, index='', limit=0,
  options = _options_with_prefix(database.name)
  transaction = self._make_txn_selector()
 
- iterator = api.streaming_read(
+ restart = functools.partial(
+ api.streaming_read,
  self._session.name, table, columns, keyset.to_pb(),
  transaction=transaction, index=index, limit=limit,
- resume_token=resume_token, options=options)
+ options=options)
+
+ iterator = _restart_on_unavailable(restart)
 
  self._read_request_count += 1
 
@@ -101,8 +133,7 @@ def read(self, table, columns, keyset, index='', limit=0,
  else:
  return StreamedResultSet(iterator)
 
- def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
- resume_token=b''):
+ def execute_sql(self, sql, params=None, param_types=None, query_mode=None):
  """Perform an ``ExecuteStreamingSql`` API request for rows in a table.
 
  :type sql: str
@@ -122,9 +153,6 @@ def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
  :param query_mode: Mode governing return of results / query plan. See
  https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode1
 
- :type resume_token: bytes
- :param resume_token: token for resuming previously-interrupted query
-
  :rtype: :class:`~google.cloud.spanner.streamed.StreamedResultSet`
  :returns: a result set instance which can be used to consume rows.
  :raises ValueError:
@@ -150,10 +178,14 @@ def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
  options = _options_with_prefix(database.name)
  transaction = self._make_txn_selector()
  api = database.spanner_api
- iterator = api.execute_streaming_sql(
+
+ restart = functools.partial(
+ api.execute_streaming_sql,
  self._session.name, sql,
  transaction=transaction, params=params_pb, param_types=param_types,
- query_mode=query_mode, resume_token=resume_token, options=options)
+ query_mode=query_mode, options=options)
+
+ iterator = _restart_on_unavailable(restart)
 
  self._read_request_count += 1
 

spanner/google/cloud/spanner/streamed.py

@@ -43,7 +43,6 @@ def __init__(self, response_iterator, source=None):
  self._counter = 0 # Counter for processed responses
  self._metadata = None # Until set from first PRS
  self._stats = None # Until set from last PRS
- self._resume_token = None # To resume from last received PRS
  self._current_row = [] # Accumulated values for incomplete row
  self._pending_chunk = None # Incomplete value
  self._source = source # Source snapshot
@@ -85,15 +84,6 @@ def stats(self):
  """
  return self._stats
 
- @property
- def resume_token(self):
- """Token for resuming interrupted read / query.
-
- :rtype: bytes
- :returns: token from last chunk of results.
- """
- return self._resume_token
-
  def _merge_chunk(self, value):
  """Merge pending chunk with next value.
 
@@ -132,7 +122,6 @@ def consume_next(self):
  """
  response = six.next(self._response_iterator)
  self._counter += 1
- self._resume_token = response.resume_token
 
  if self._metadata is None: # first response
  metadata = self._metadata = response.metadata

spanner/tests/unit/test_session.py

@@ -265,7 +265,6 @@ def test_read(self):
  KEYSET = KeySet(keys=KEYS)
  INDEX = 'email-address-index'
  LIMIT = 20
- TOKEN = b'DEADBEEF'
  database = _Database(self.DATABASE_NAME)
  session = self._make_one(database)
  session._session_id = 'DEADBEEF'
@@ -279,28 +278,26 @@ def __init__(self, session, **kwargs):
  self._session = session
  self._kwargs = kwargs.copy()
 
- def read(self, table, columns, keyset, index='', limit=0,
- resume_token=b''):
+ def read(self, table, columns, keyset, index='', limit=0):
  _read_with.append(
- (table, columns, keyset, index, limit, resume_token))
+ (table, columns, keyset, index, limit))
  return expected
 
  with _Monkey(MUT, Snapshot=_Snapshot):
  found = session.read(
  TABLE_NAME, COLUMNS, KEYSET,
- index=INDEX, limit=LIMIT, resume_token=TOKEN)
+ index=INDEX, limit=LIMIT)
 
  self.assertIs(found, expected)
 
  self.assertEqual(len(_read_with), 1)
- (table, columns, key_set, index, limit, resume_token) = _read_with[0]
+ (table, columns, key_set, index, limit) = _read_with[0]
 
  self.assertEqual(table, TABLE_NAME)
  self.assertEqual(columns, COLUMNS)
  self.assertEqual(key_set, KEYSET)
  self.assertEqual(index, INDEX)
  self.assertEqual(limit, LIMIT)
- self.assertEqual(resume_token, TOKEN)
 
  def test_execute_sql_not_created(self):
  SQL = 'SELECT first_name, age FROM citizens'
@@ -330,25 +327,23 @@ def __init__(self, session, **kwargs):
  self._kwargs = kwargs.copy()
 
  def execute_sql(
- self, sql, params=None, param_types=None, query_mode=None,
- resume_token=None):
+ self, sql, params=None, param_types=None, query_mode=None):
  _executed_sql_with.append(
- (sql, params, param_types, query_mode, resume_token))
+ (sql, params, param_types, query_mode))
  return expected
 
  with _Monkey(MUT, Snapshot=_Snapshot):
- found = session.execute_sql(SQL, resume_token=TOKEN)
+ found = session.execute_sql(SQL)
 
  self.assertIs(found, expected)
 
  self.assertEqual(len(_executed_sql_with), 1)
- sql, params, param_types, query_mode, token = _executed_sql_with[0]
+ sql, params, param_types, query_mode = _executed_sql_with[0]
 
  self.assertEqual(sql, SQL)
  self.assertEqual(params, None)
  self.assertEqual(param_types, None)
  self.assertEqual(query_mode, None)
- self.assertEqual(token, TOKEN)
 
  def test_batch_not_created(self):
  database = _Database(self.DATABASE_NAME)