Repository: googleapis/google-cloud-python

Issue:
Cloud spanner client does not seem to be retrying broken streams
[Documentation](http://google-cloud-python.readthedocs.io/en/latest/spanner/snapshot-usage.html) for snapshot transactions mentions:

```
If streaming a chunk fails due to a “resumable” error, Session.read() retries the StreamingRead API reqeust, passing the resume_token from the last partial result streamed.
```

This is the correct thing to do but the client actually does not seem to be doing this. I looked around a little bit in the code and I could find anyplace where it is retrying the read with the resume token. 
This is the place where the stream would throw an error if it broke:
https://github.com/GoogleCloudPlatform/google-cloud-python/blob/9d9b6c0708ae6553458ad6107e5ad8efb23762e8/spanner/google/cloud/spanner/streamed.py#L132

 I dont see any exception handling happening here and it looks like that the error would just be propagated back to the user.


Discussion:
cc @bjwatson @tseaver @lukesneeringer 
Our retry logic is in GAX, but leaving this open for @tseaver to confirm that it is applicable here.
This needs to be done at a higher level than GAX, since we send resume tokens along with the responses. So if a stream breaks, instead of retrying the whole stream, the library should just pass the resume token back in the request which will cause the stream to be resumed from that point onwards. It is further complicated since cloud spanner might not send resume tokens with all responses so the client library needs to buffer responses till it sees one with resume token and only then yield them to the user. This will allow it to safely resume in case the stream breaks.
This is the Java implementation of the same logic:
https://github.com/GoogleCloudPlatform/google-cloud-java/blob/master/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java#L2166

What does GAX do for streaming calls? Does it retry the whole stream if it fails? What if the stream had yielded all but the last 1 response before failing? In that case retrying the whole stream is very wasteful. If GAX does do retries for streaming calls, we need to disable them.
> What does GAX do for streaming calls? Does it retry the whole stream if it fails? What if the stream had yielded all but the last 1 response before failing? In that case retrying the whole stream is very wasteful. If GAX does do retries for streaming calls, we need to disable them.

gax is pretty naive in what it does for anything other than simple RPCs. I'm writing a doc internally on this now. :)
@lukesneeringer @vkedia Does this need to be Beta blocking?
Till this is resolved can we fix the documentation to indicate that users would need to resume the stream themselves by passing back the resume token if available?
`Snapshot.read` and `Snapshot.execute_sql` return an instance of `StreamedResultSet`.  That instance does not contain the information needed to reconstruct the original call inside its `consume_next`, which means that we cannot make handling such errors seamless to the user out-of-the box.

We could attempt to add enough information to the `StreamedResultSet` (a curried partial, maybe?) to allow such retries.  E.g., from `read()`:

```python
    retry = functools.partial(self.read, table, columns, keyset, index, limit)
    return StreamedResultSet(iterator, retry=retry)
```

or from `execute_sql`:

```python
    retry = functools.partial(self.execute_sql, params, param_types, query_mode)
    return StreamedResultSet(iterator, retry=retry)
```

The only other way I can see to hide such errors from the user would be to add callback-based methods which took the arguments to be passed to `read` / `execute_sql`, executed them, and then iterated the result set, invoking the callback for each row successfully pulled from the stream.

E.g.:

```python
    def handle_row(row):
        ...

    snapshot.read_cb(table, keyset, ..., callback=handle_row)

    snapshot.execute_sql(sql, ..., callback=handle_row)
```
I don't understand how the callback solution fixes this. Can you please elaborate? How would you retry in that case?
Your option 1 is similar to what we do in Java as well. Note that when you retry you will have to pass back the resume_token, so the retry function could possible take a resume token as an argument.

Also note this is more complex because we do not guarantee that all PartialResultSet would have a resume token. So if you yield a Row to the user but the PRS that contained that Row did not have a resume_token, you cannot retry safely. So what you will need to do is store PRS's in a buffer till you see one with resume_token. Only then yield them to the user. So essentially you guarantee that every time you yield something to the user it is at a resumable boundary.
The curried partial is the pattern we'll use in `google.api.core`. Instead of passing the curried partial to `StreamedResultSet`, can you construct an intermediary iterator wrapper for the request?

```python
class _RetryingStreamIterator(object):
    def __init__(self, iter, request):
        self._iter = iter
        self._request = request
        self._resume_token = None

    def __iter__(self):
        return self

    def __next__(self):
        try:
            next = six.next(self._iter)
            self._resume_token = next.resumeToken
        except grpc.RpcError as exc:
            if (exc.status_code == grpc.StatusCode.RESUMABLE
                    and self.resumeToken is not None):
                self._iter = self._request(resume_token=self._resume_token)
                return six.next(self._iter)
            else:
                raise

...


# read/execute_sql:

request = 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)

iterator = _RetryingStreamIterator(request(), request)

if self._multi_use:
    return StreamedResultSet(iterator, source=self)
else:
    return StreamedResultSet(iterator)
```

I applied the `release blocking` label temporarily, to make sure we reach some kind of resolution on this before Beta. Either there's a quick and easy fix, or we update the docs so that we're not misleading users. In the latter case, we'll want to fix this shortly after declaring Beta.
@bjwatson I've removed the "release blocking" label becuase #3790 documents the current, application-level retry semantics.
Great, thanks @tseaver!
I left a comment on that PR. The documentation is not completely correct and as I mention their it is really tricky to tell the user what they need to do in case of a broken stream (which is why we handle that in the library). So I think it is better to fix this.
I thought about this a bit more and it is really hard for users to build their own retry logic around the current API and will be error prone. It is better to just fix it in the client code. 
@bjwatson Can you add back the release blocking label to this?
/cc @jonparrott Who has been working on the retry design
If it is release blocking, the _minimum_ that this will push us back is two weeks, since @tseaver will be mostly out.
I see. I also realized now that this might be a breaking change as well. That is because once we implement retries in the client, there is no reason for us to expose the concept of resume_token to the users. It will just be confusing. So we should remove the `resume_token` [method](https://github.com/GoogleCloudPlatform/google-cloud-python/blob/master/spanner/google/cloud/spanner/streamed.py#L89) from StreamedIterator. 
We should also remove `resume_token` as an [argument](https://github.com/GoogleCloudPlatform/google-cloud-python/blob/master/spanner/google/cloud/spanner/snapshot.py#L72) to read/query.
With #3819 in we have the core retry functionality needed for this. I do not plan on adding a general-purpose iterator retry decorator to api.core at the moment.
Looking at using the retry features from #3819 to implement this.  @vkedia, is there a specification for the set of "retryable" exceptions which would be raised during streaming iteration?
Also, is there a spanner-specific deadline after which retries should abort?
@jonparrott Are we doing anything about "sensible defaults" for the `initial` and `maximum` arguments to `google.api.core.retry.exponential_sleep_generator`?
@tseaver the defaults to the high-level `Retry` are sensible. The expectation for gapic clients is that they'll populate those values based on the gapic client config json file. I haven't yet sent a PR for that code, but here's a [gist](https://gist.github.com/jonparrott/afbb97a9fc273d306d0d19feeb99aec6) from my prototype.

@lukesneeringer can help with the gapic config part.
One possibly valid approach in the near term is to use `Retry`'s defaults until I'm done replacing gax, in which case using the gapic config should be trivial. WDYT, @lukesneeringer?
@jonparrott I missed seeing the `Retry` class, which looks to have been added after #3819 (update, added in #3835).

Also, I'm unsure how GAPIC config fits in with retries for the streaming-result-set iteration, where code in this library is responsible for tracking / passing back the `resume_token` to "restart" the result set after a transient failure.
> Also, I'm unsure how GAPIC config fits in with retries for the streaming-result-set iteration, where code in this library is responsible for tracking / passing back the resume_token to "restart" the result set after a transient failure.

Yeah in general I wouldn't sweat that for your purposes here- sorry to mislead you on that (it's monday..)

Basically, two levels of retries will be involved here: the "iterator" retry and the "method" retry. e.g.: calling `next(iter)` will catch errors and and retry calling the *method* to re-establish the stream. The method itself is wrapped with its own retry. 

It's turtles all the way down.

I don't think you need to work about the gapic config for the iterator retry- you can just use the default retry or customize as you feel is appropriate.
So I need to know what are the"transient" errors which may be propagated from the "iterator" retry, and which the "method" (`StreamedResultSet.consume_next`) will need to trap in order to restart the API, passing the previously-captured `resume_token`.  I'm also unclear whether there is supposed to be any backoff / jitter assocated with the "restart" (vs. invisible-to-me retries of the "iterator" API calls).
Per @blowmage:  the [Ruby implementation](https://github.com/GoogleCloudPlatform/google-cloud-ruby/blob/master/google-cloud-spanner/lib/google/cloud/spanner/results.rb#L143) appears to only do the "restart" for `UNAVAILABLE`.
Yes you should try only for UNAVAILABLE with exponential backoff.

On Mon, Aug 28, 2017 at 9:57 AM, Tres Seaver <notifications@github.com>
wrote:

> Per @blowmage <https://github.com/blowmage>: the Ruby implementation
> <https://github.com/GoogleCloudPlatform/google-cloud-ruby/blob/master/google-cloud-spanner/lib/google/cloud/spanner/results.rb#L143>
> appears to only do the "restart" for UNAVAILABLE.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/GoogleCloudPlatform/google-cloud-python/issues/3775#issuecomment-325411015>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ATdef5aCSTCk6do3cbvM_cAhx4zqEl6Xks5scvFjgaJpZM4OyakU>
> .
>

@jonparrott the "stock" bits in `google.api.core.retry` don't really fit this case:  I need to inject a "fixup" when the appropriate exception (`UNAVAILABLE`) is raised, *before* retrying the stream iteration.  Something like:

```python

    def consume_next(self):
        """Consume the next partial result set from the stream.

        Parse the result set into new/existing rows in :attr:`_rows`
        """
        sleep_generator = retry.exponential_sleep_generator(
            initial=retry._DEFAULT_INITIAL_DELAY,
            maximum=retry._DEFAULT_MAXIMUM_DELAY)

        for sleep in sleep_generator:
            try:
                response = six.next(self._response_iterator)
            except exceptions.ServiceUnavailable:

                if self._resume_token in (None, b''):
                    raise

                self._response_iterator = self._retry(
                    resume_token=self._resume_token)
            else:
                break
```
@tseaver that seems pretty reasonable. (it's one of the reasons I made exponential_sleep_generator public)

You should set a deadline of some sort, though.
Does the Python client allow users to specify deadline? If it does then the
same deadline should be set adjusting for the time time elapsed so far.


On Mon, Aug 28, 2017 at 10:43 AM, Jon Wayne Parrott <
notifications@github.com> wrote:

> @tseaver <https://github.com/tseaver> that seems pretty reasonable. (it's
> one of the reasons I made exponential_sleep_generator public)
>
> You should set a deadline of some sort, though.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/GoogleCloudPlatform/google-cloud-python/issues/3775#issuecomment-325423984>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ATdef1-qD59qWEW1G7iw2pW8ElDJyjstks5scvxNgaJpZM4OyakU>
> .
>

And of course I have to add my own deadline handling.

BTW, why does `google.api.core.retry.retry_target` use the more convoluted `datetime` stuff to compute the deadline, rather than just `time.time` plus a constant number of seconds?
> BTW, why does google.api.core.retry.retry_target use the more convoluted datetime stuff to compute the deadline, rather than just time.time plus a constant number of seconds?

No particular reason other than being functionally equivalent and having on single source of truth for time in tests (`datetime_helpers.utcnow()`).

> And of course I have to add my own deadline handling.

How so?
@jonparrott:

> >  And of course I have to add my own deadline handling.

> How so?

As you pointed out above, I'm doing the low level work myself.

ISTM that the `retry_target` (and therefore `Retry`) could grow some feature to signal a "fixup" to happen between the retries, to accommodate this kind of stateful behavior.
I'm still not sure what you mean? Just pass the user's deadline into retry_target?
I can't use `retry_target`:  it doesn't let me issue the "restart" request and capture the new iterator.  And so I need to do my own deadline handling.
Thanks for the PR - I think I see now, so it seems like if you did this:

```python
retry_target(functools.partial(six.next, self._response_iterator))
```

It wouldn't let you catch the error and replace `self._response_iterator`.

I could imagine a hack using the predicate and some getattr magic, but that sounds gross to me.

Is there a minimal change you could make to `retry_target` to make it work for you- perhaps an optional `on_exception` callback? If so, I'm happy to review and merge that.
Also - I've thought about a `retry_iterator` function - maybe it makes more sense to just go ahead and add that to core?

Resolving patch:
diff --git a/spanner/google/cloud/spanner/session.py b/spanner/google/cloud/spanner/session.py
--- a/spanner/google/cloud/spanner/session.py
+++ b/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.
diff --git a/spanner/google/cloud/spanner/snapshot.py b/spanner/google/cloud/spanner/snapshot.py
--- a/spanner/google/cloud/spanner/snapshot.py
+++ b/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
 
diff --git a/spanner/google/cloud/spanner/streamed.py b/spanner/google/cloud/spanner/streamed.py
--- a/spanner/google/cloud/spanner/streamed.py
+++ b/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


Test patch:
diff --git a/spanner/tests/unit/test_session.py b/spanner/tests/unit/test_session.py
--- a/spanner/tests/unit/test_session.py
+++ b/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)
diff --git a/spanner/tests/unit/test_snapshot.py b/spanner/tests/unit/test_snapshot.py
--- a/spanner/tests/unit/test_snapshot.py
+++ b/spanner/tests/unit/test_snapshot.py
@@ -15,6 +15,8 @@
 
 import unittest
 
+import mock
+
 from google.cloud._testing import _GAXBaseAPI
 
 
@@ -31,6 +33,85 @@
 PARAMS_WITH_BYTES = {'bytes': b'DEADBEEF'}
 
 
+class Test_restart_on_unavailable(unittest.TestCase):
+
+    def _call_fut(self, restart):
+        from google.cloud.spanner.snapshot import _restart_on_unavailable
+
+        return _restart_on_unavailable(restart)
+
+    def _make_item(self, value, resume_token=''):
+        return mock.Mock(
+            value=value, resume_token=resume_token,
+            spec=['value', 'resume_token'])
+
+    def test_iteration_w_empty_raw(self):
+        ITEMS = ()
+        raw = _MockIterator()
+        restart = mock.Mock(spec=[], return_value=raw)
+        resumable = self._call_fut(restart)
+        self.assertEqual(list(resumable), [])
+
+    def test_iteration_w_non_empty_raw(self):
+        ITEMS = (self._make_item(0), self._make_item(1))
+        raw = _MockIterator(*ITEMS)
+        restart = mock.Mock(spec=[], return_value=raw)
+        resumable = self._call_fut(restart)
+        self.assertEqual(list(resumable), list(ITEMS))
+        restart.assert_called_once_with()
+
+    def test_iteration_w_raw_w_resume_tken(self):
+        ITEMS = (
+            self._make_item(0),
+            self._make_item(1, resume_token='DEADBEEF'),
+            self._make_item(2),
+            self._make_item(3),
+        )
+        raw = _MockIterator(*ITEMS)
+        restart = mock.Mock(spec=[], return_value=raw)
+        resumable = self._call_fut(restart)
+        self.assertEqual(list(resumable), list(ITEMS))
+        restart.assert_called_once_with()
+
+    def test_iteration_w_raw_raising_unavailable(self):
+        FIRST = (
+            self._make_item(0),
+            self._make_item(1, resume_token='DEADBEEF'),
+        )
+        SECOND = (  # discarded after 503
+            self._make_item(2),
+        )
+        LAST = (
+            self._make_item(3),
+        )
+        before = _MockIterator(*(FIRST + SECOND), fail_after=True)
+        after = _MockIterator(*LAST)
+        restart = mock.Mock(spec=[], side_effect=[before, after])
+        resumable = self._call_fut(restart)
+        self.assertEqual(list(resumable), list(FIRST + LAST))
+        self.assertEqual(
+            restart.mock_calls,
+            [mock.call(), mock.call(resume_token='DEADBEEF')])
+
+    def test_iteration_w_raw_raising_unavailable_after_token(self):
+        FIRST = (
+            self._make_item(0),
+            self._make_item(1, resume_token='DEADBEEF'),
+        )
+        SECOND = (
+            self._make_item(2),
+            self._make_item(3),
+        )
+        before = _MockIterator(*FIRST, fail_after=True)
+        after = _MockIterator(*SECOND)
+        restart = mock.Mock(spec=[], side_effect=[before, after])
+        resumable = self._call_fut(restart)
+        self.assertEqual(list(resumable), list(FIRST + SECOND))
+        self.assertEqual(
+            restart.mock_calls,
+            [mock.call(), mock.call(resume_token='DEADBEEF')])
+
+
 class Test_SnapshotBase(unittest.TestCase):
 
     PROJECT_ID = 'project-id'
@@ -95,7 +176,7 @@ def test_read_grpc_error(self):
         derived = self._makeDerived(session)
 
         with self.assertRaises(GaxError):
-            derived.read(TABLE_NAME, COLUMNS, KEYSET)
+            list(derived.read(TABLE_NAME, COLUMNS, KEYSET))
 
         (r_session, table, columns, key_set, transaction, index,
          limit, resume_token, options) = api._streaming_read_with
@@ -152,7 +233,7 @@ def _read_helper(self, multi_use, first=True, count=0):
         TOKEN = b'DEADBEEF'
         database = _Database()
         api = database.spanner_api = _FauxSpannerAPI(
-            _streaming_read_response=_MockCancellableIterator(*result_sets))
+            _streaming_read_response=_MockIterator(*result_sets))
         session = _Session(database)
         derived = self._makeDerived(session)
         derived._multi_use = multi_use
@@ -162,7 +243,7 @@ def _read_helper(self, multi_use, first=True, count=0):
 
         result_set = derived.read(
             TABLE_NAME, COLUMNS, KEYSET,
-            index=INDEX, limit=LIMIT, resume_token=TOKEN)
+            index=INDEX, limit=LIMIT)
 
         self.assertEqual(derived._read_request_count, count + 1)
 
@@ -172,6 +253,7 @@ def _read_helper(self, multi_use, first=True, count=0):
             self.assertIsNone(result_set._source)
 
         result_set.consume_all()
+
         self.assertEqual(list(result_set.rows), VALUES)
         self.assertEqual(result_set.metadata, metadata_pb)
         self.assertEqual(result_set.stats, stats_pb)
@@ -193,7 +275,7 @@ def _read_helper(self, multi_use, first=True, count=0):
             self.assertTrue(transaction.single_use.read_only.strong)
         self.assertEqual(index, INDEX)
         self.assertEqual(limit, LIMIT)
-        self.assertEqual(resume_token, TOKEN)
+        self.assertEqual(resume_token, b'')
         self.assertEqual(options.kwargs['metadata'],
                          [('google-cloud-resource-prefix', database.name)])
 
@@ -229,7 +311,7 @@ def test_execute_sql_grpc_error(self):
         derived = self._makeDerived(session)
 
         with self.assertRaises(GaxError):
-            derived.execute_sql(SQL_QUERY)
+            list(derived.execute_sql(SQL_QUERY))
 
         (r_session, sql, transaction, params, param_types,
          resume_token, query_mode, options) = api._executed_streaming_sql_with
@@ -288,7 +370,7 @@ def _execute_sql_helper(self, multi_use, first=True, count=0):
             PartialResultSet(values=VALUE_PBS[0], metadata=metadata_pb),
             PartialResultSet(values=VALUE_PBS[1], stats=stats_pb),
         ]
-        iterator = _MockCancellableIterator(*result_sets)
+        iterator = _MockIterator(*result_sets)
         database = _Database()
         api = database.spanner_api = _FauxSpannerAPI(
             _execute_streaming_sql_response=iterator)
@@ -301,7 +383,7 @@ def _execute_sql_helper(self, multi_use, first=True, count=0):
 
         result_set = derived.execute_sql(
             SQL_QUERY_WITH_PARAM, PARAMS, PARAM_TYPES,
-            query_mode=MODE, resume_token=TOKEN)
+            query_mode=MODE)
 
         self.assertEqual(derived._read_request_count, count + 1)
 
@@ -311,6 +393,7 @@ def _execute_sql_helper(self, multi_use, first=True, count=0):
             self.assertIsNone(result_set._source)
 
         result_set.consume_all()
+
         self.assertEqual(list(result_set.rows), VALUES)
         self.assertEqual(result_set.metadata, metadata_pb)
         self.assertEqual(result_set.stats, stats_pb)
@@ -333,7 +416,7 @@ def _execute_sql_helper(self, multi_use, first=True, count=0):
         self.assertEqual(params, expected_params)
         self.assertEqual(param_types, PARAM_TYPES)
         self.assertEqual(query_mode, MODE)
-        self.assertEqual(resume_token, TOKEN)
+        self.assertEqual(resume_token, b'')
         self.assertEqual(options.kwargs['metadata'],
                          [('google-cloud-resource-prefix', database.name)])
 
@@ -358,20 +441,6 @@ def test_execute_sql_w_multi_use_w_first_w_count_gt_0(self):
             self._execute_sql_helper(multi_use=True, first=True, count=1)
 
 
-class _MockCancellableIterator(object):
-
-    cancel_calls = 0
-
-    def __init__(self, *values):
-        self.iter_values = iter(values)
-
-    def next(self):
-        return next(self.iter_values)
-
-    def __next__(self):  # pragma: NO COVER Py3k
-        return self.next()
-
-
 class TestSnapshot(unittest.TestCase):
 
     PROJECT_ID = 'project-id'
@@ -725,7 +794,7 @@ def begin_transaction(self, session, options_, options=None):
     # pylint: disable=too-many-arguments
     def streaming_read(self, session, table, columns, key_set,
                        transaction=None, index='', limit=0,
-                       resume_token='', options=None):
+                       resume_token=b'', options=None):
         from google.gax.errors import GaxError
 
         self._streaming_read_with = (
@@ -738,7 +807,7 @@ def streaming_read(self, session, table, columns, key_set,
 
     def execute_streaming_sql(self, session, sql, transaction=None,
                               params=None, param_types=None,
-                              resume_token='', query_mode=None, options=None):
+                              resume_token=b'', query_mode=None, options=None):
         from google.gax.errors import GaxError
 
         self._executed_streaming_sql_with = (
@@ -747,3 +816,25 @@ def execute_streaming_sql(self, session, sql, transaction=None,
         if self._random_gax_error:
             raise GaxError('error')
         return self._execute_streaming_sql_response
+
+
+class _MockIterator(object):
+
+    def __init__(self, *values, **kw):
+        self._iter_values = iter(values)
+        self._fail_after = kw.pop('fail_after', False)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        from google.api.core.exceptions import ServiceUnavailable
+
+        try:
+            return next(self._iter_values)
+        except StopIteration:
+            if self._fail_after:
+                raise ServiceUnavailable('testing')
+            raise
+
+    next = __next__
diff --git a/spanner/tests/unit/test_streamed.py b/spanner/tests/unit/test_streamed.py
--- a/spanner/tests/unit/test_streamed.py
+++ b/spanner/tests/unit/test_streamed.py
@@ -36,7 +36,6 @@ def test_ctor_defaults(self):
         self.assertEqual(streamed.rows, [])
         self.assertIsNone(streamed.metadata)
         self.assertIsNone(streamed.stats)
-        self.assertIsNone(streamed.resume_token)
 
     def test_ctor_w_source(self):
         iterator = _MockCancellableIterator()
@@ -47,7 +46,6 @@ def test_ctor_w_source(self):
         self.assertEqual(streamed.rows, [])
         self.assertIsNone(streamed.metadata)
         self.assertIsNone(streamed.stats)
-        self.assertIsNone(streamed.resume_token)
 
     def test_fields_unset(self):
         iterator = _MockCancellableIterator()
@@ -669,7 +667,6 @@ def test_consume_next_first_set_partial(self):
         self.assertEqual(streamed.rows, [])
         self.assertEqual(streamed._current_row, BARE)
         self.assertEqual(streamed.metadata, metadata)
-        self.assertEqual(streamed.resume_token, result_set.resume_token)
         self.assertEqual(source._transaction_id, TXN_ID)
 
     def test_consume_next_first_set_partial_existing_txn_id(self):
@@ -691,7 +688,6 @@ def test_consume_next_first_set_partial_existing_txn_id(self):
         self.assertEqual(streamed.rows, [])
         self.assertEqual(streamed._current_row, BARE)
         self.assertEqual(streamed.metadata, metadata)
-        self.assertEqual(streamed.resume_token, result_set.resume_token)
         self.assertEqual(source._transaction_id, TXN_ID)
 
     def test_consume_next_w_partial_result(self):
@@ -711,7 +707,6 @@ def test_consume_next_w_partial_result(self):
         self.assertEqual(streamed.rows, [])
         self.assertEqual(streamed._current_row, [])
         self.assertEqual(streamed._pending_chunk, VALUES[0])
-        self.assertEqual(streamed.resume_token, result_set.resume_token)
 
     def test_consume_next_w_pending_chunk(self):
         FIELDS = [
@@ -737,7 +732,6 @@ def test_consume_next_w_pending_chunk(self):
         ])
         self.assertEqual(streamed._current_row, [BARE[6]])
         self.assertIsNone(streamed._pending_chunk)
-        self.assertEqual(streamed.resume_token, result_set.resume_token)
 
     def test_consume_next_last_set(self):
         FIELDS = [
@@ -761,7 +755,6 @@ def test_consume_next_last_set(self):
         self.assertEqual(streamed.rows, [BARE])
         self.assertEqual(streamed._current_row, [])
         self.assertEqual(streamed._stats, stats)
-        self.assertEqual(streamed.resume_token, result_set.resume_token)
 
     def test_consume_all_empty(self):
         iterator = _MockCancellableIterator()