diff --git a/api_core/google/api_core/gapic_v1/method.py b/api_core/google/api_core/gapic_v1/method.py index 3a689afbd137..425e81e5c139 100644 --- a/api_core/google/api_core/gapic_v1/method.py +++ b/api_core/google/api_core/gapic_v1/method.py @@ -18,13 +18,8 @@ pagination, and long-running operations to gRPC methods. """ -import functools - -import six - from google.api_core import general_helpers from google.api_core import grpc_helpers -from google.api_core import page_iterator from google.api_core import timeout from google.api_core.gapic_v1 import client_info @@ -232,38 +227,3 @@ def get_topic(name, timeout=None): _GapicCallable( func, default_retry, default_timeout, user_agent_metadata=user_agent_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/api_core/google/api_core/operations_v1/operations_client.py b/api_core/google/api_core/operations_v1/operations_client.py index 93a823f05d6d..c8eff4879104 100644 --- a/api_core/google/api_core/operations_v1/operations_client.py +++ b/api_core/google/api_core/operations_v1/operations_client.py @@ -35,7 +35,10 @@ /operations.proto """ +import functools + from google.api_core import gapic_v1 +from google.api_core import page_iterator from google.api_core.operations_v1 import operations_client_config from google.longrunning import operations_pb2 @@ -72,12 +75,6 @@ def __init__(self, channel, client_config=operations_client_config.config): default_retry=method_configs['ListOperations'].retry, default_timeout=method_configs['ListOperations'].timeout) - self._list_operations = gapic_v1.method.wrap_with_paging( - self._list_operations, - 'operations', - 'page_token', - 'next_page_token') - self._cancel_operation = gapic_v1.method.wrap_method( self.operations_stub.CancelOperation, default_retry=method_configs['CancelOperation'].retry, @@ -182,7 +179,20 @@ def list_operations( # Create the request object. request = operations_pb2.ListOperationsRequest( name=name, filter=filter_) - return self._list_operations(request, retry=retry, timeout=timeout) + + # Create the method used to fetch pages + method = functools.partial( + self._list_operations, retry=retry, timeout=timeout) + + iterator = page_iterator.GRPCIterator( + client=None, + method=method, + request=request, + items_field='operations', + request_token_field='page_token', + response_token_field='next_page_token') + + return iterator def cancel_operation( self, name, diff --git a/api_core/tests/unit/gapic/test_method.py b/api_core/tests/unit/gapic/test_method.py index 281463eb6c12..0c41f87d8c36 100644 --- a/api_core/tests/unit/gapic/test_method.py +++ b/api_core/tests/unit/gapic/test_method.py @@ -179,38 +179,3 @@ 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__', '__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') - - 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'