From 9ad88727df113ec0114c4283612617a23e38baa5 Mon Sep 17 00:00:00 2001 From: Jon Wayne Parrott Date: Tue, 26 Sep 2017 12:57:01 -0700 Subject: [PATCH 1/2] Add wrap_with_paging --- core/google/api/core/gapic_v1/method.py | 41 ++++++++++++++++++- core/tests/unit/api_core/gapic/test_method.py | 34 +++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/core/google/api/core/gapic_v1/method.py b/core/google/api/core/gapic_v1/method.py index 8c92200f9252..ee7da0a26959 100644 --- a/core/google/api/core/gapic_v1/method.py +++ b/core/google/api/core/gapic_v1/method.py @@ -18,10 +18,13 @@ pagination, and long-running operations to gRPC methods. """ +import functools import platform import pkg_resources +import six +from google.api.core import page_iterator from google.api.core import timeout from google.api.core.helpers import grpc_helpers @@ -238,4 +241,40 @@ def get_topic(name, timeout=None): if metadata is not None: metadata = _prepare_metadata(metadata) - return _GapicCallable(func, default_retry, default_timeout, metadata) + return six.wraps(func)( + _GapicCallable(func, default_retry, default_timeout, metadata)) + + +def wrap_with_paging( + func, items_field, request_token_field, response_token_field): + """Wrap an RPC method to return a page iterator. + + Args: + func (Callable): The RPC method. This should already have been + wrapped with common functionality using :func:`wrap_method`. + request (protobuf.Message): The request message. + items_field (str): The field in the response message that has the + items for the page. + request_token_field (str): The field in the request message used to + specify the page token. + response_token_field (str): The field in the response message that has + the token for the next page. + + Returns: + Callable: Returns a callable that when invoked will call the RPC + method and return a + :class:`google.api.core.page_iterator.Iterator`. + """ + @six.wraps(func) + def paged_method(request, **kwargs): + """Wrapper that invokes a method and returns a page iterator.""" + iterator = page_iterator.GRPCIterator( + client=None, + method=functools.partial(func, **kwargs), + request=request, + items_field=items_field, + request_token_field=request_token_field, + response_token_field=response_token_field) + return iterator + + return paged_method diff --git a/core/tests/unit/api_core/gapic/test_method.py b/core/tests/unit/api_core/gapic/test_method.py index f2e5abb17407..c3fea4c860f7 100644 --- a/core/tests/unit/api_core/gapic/test_method.py +++ b/core/tests/unit/api_core/gapic/test_method.py @@ -18,6 +18,7 @@ from google.api.core import retry from google.api.core import timeout import google.api.core.gapic_v1.method +import google.api.core.page_iterator def test_wrap_method_basic(): @@ -144,3 +145,36 @@ def test_wrap_method_with_overriding_timeout_as_a_number(): assert result == 42 method.assert_called_once_with(timeout=22, metadata=mock.ANY) + + +def test_wrap_with_paging(): + page_one = mock.Mock( + spec=['items', 'page_token', 'next_page_token'], + items=[1, 2], + next_page_token='icanhasnextpls') + page_two = mock.Mock( + spec=['items', 'page_token', 'next_page_token'], + items=[3, 4], + next_page_token=None) + method = mock.Mock(spec=['__call__'], side_effect=(page_one, page_two)) + + wrapped_method = google.api.core.gapic_v1.method.wrap_with_paging( + method, 'items', 'page_token', 'next_page_token') + + request = mock.Mock(spec=['page_token'], page_token=None) + result = wrapped_method(request, extra='param') + + # Should return an iterator and should not have actually called the + # method yet. + assert isinstance(result, google.api.core.page_iterator.Iterator) + method.assert_not_called() + assert request.page_token is None + + # Draining the iterator should call the method until no more pages are + # returned. + results = list(result) + + assert results == [1, 2, 3, 4] + assert method.call_count == 2 + method.assert_called_with(request, extra='param') + assert request.page_token == 'icanhasnextpls' From 514faf30c004ed6bfb18a8ffad699969f534f616 Mon Sep 17 00:00:00 2001 From: Jon Wayne Parrott Date: Tue, 26 Sep 2017 13:04:33 -0700 Subject: [PATCH 2/2] Fixup for py27 --- core/tests/unit/api_core/gapic/test_method.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/core/tests/unit/api_core/gapic/test_method.py b/core/tests/unit/api_core/gapic/test_method.py index c3fea4c860f7..c513940bdf4b 100644 --- a/core/tests/unit/api_core/gapic/test_method.py +++ b/core/tests/unit/api_core/gapic/test_method.py @@ -156,7 +156,9 @@ def test_wrap_with_paging(): spec=['items', 'page_token', 'next_page_token'], items=[3, 4], next_page_token=None) - method = mock.Mock(spec=['__call__'], side_effect=(page_one, page_two)) + method = mock.Mock( + spec=['__call__', '__name__'], side_effect=(page_one, page_two)) + method.__name__ = 'mockmethod' wrapped_method = google.api.core.gapic_v1.method.wrap_with_paging( method, 'items', 'page_token', 'next_page_token')