1#!/usr/bin/env python 2# 3# Copyright 2015 Google Inc. 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16 17"""A helper function that executes a series of List queries for many APIs.""" 18 19from apitools.base.py import encoding 20import six 21 22__all__ = [ 23 'YieldFromList', 24] 25 26 27def _GetattrNested(message, attribute): 28 """Gets a possibly nested attribute. 29 30 Same as getattr() if attribute is a string; 31 if attribute is a tuple, returns the nested attribute referred to by 32 the fields in the tuple as if they were a dotted accessor path. 33 34 (ex _GetattrNested(msg, ('foo', 'bar', 'baz')) gets msg.foo.bar.baz 35 """ 36 if isinstance(attribute, six.string_types): 37 return getattr(message, attribute) 38 elif len(attribute) == 0: 39 return message 40 else: 41 return _GetattrNested(getattr(message, attribute[0]), attribute[1:]) 42 43 44def _SetattrNested(message, attribute, value): 45 """Sets a possibly nested attribute. 46 47 Same as setattr() if attribute is a string; 48 if attribute is a tuple, sets the nested attribute referred to by 49 the fields in the tuple as if they were a dotted accessor path. 50 51 (ex _SetattrNested(msg, ('foo', 'bar', 'baz'), 'v') sets msg.foo.bar.baz 52 to 'v' 53 """ 54 if isinstance(attribute, six.string_types): 55 return setattr(message, attribute, value) 56 elif len(attribute) < 1: 57 raise ValueError("Need an attribute to set") 58 elif len(attribute) == 1: 59 return setattr(message, attribute[0], value) 60 else: 61 return setattr(_GetattrNested(message, attribute[:-1]), 62 attribute[-1], value) 63 64 65def YieldFromList( 66 service, request, global_params=None, limit=None, batch_size=100, 67 method='List', field='items', predicate=None, 68 current_token_attribute='pageToken', 69 next_token_attribute='nextPageToken', 70 batch_size_attribute='maxResults', 71 get_field_func=_GetattrNested): 72 """Make a series of List requests, keeping track of page tokens. 73 74 Args: 75 service: apitools_base.BaseApiService, A service with a .List() method. 76 request: protorpc.messages.Message, The request message 77 corresponding to the service's .List() method, with all the 78 attributes populated except the .maxResults and .pageToken 79 attributes. 80 global_params: protorpc.messages.Message, The global query parameters to 81 provide when calling the given method. 82 limit: int, The maximum number of records to yield. None if all available 83 records should be yielded. 84 batch_size: int, The number of items to retrieve per request. 85 method: str, The name of the method used to fetch resources. 86 field: str, The field in the response that will be a list of items. 87 predicate: lambda, A function that returns true for items to be yielded. 88 current_token_attribute: str or tuple, The name of the attribute in a 89 request message holding the page token for the page being 90 requested. If a tuple, path to attribute. 91 next_token_attribute: str or tuple, The name of the attribute in a 92 response message holding the page token for the next page. If a 93 tuple, path to the attribute. 94 batch_size_attribute: str or tuple, The name of the attribute in a 95 response message holding the maximum number of results to be 96 returned. None if caller-specified batch size is unsupported. 97 If a tuple, path to the attribute. 98 get_field_func: Function that returns the items to be yielded. Argument 99 is response message, and field. 100 101 Yields: 102 protorpc.message.Message, The resources listed by the service. 103 104 """ 105 request = encoding.CopyProtoMessage(request) 106 _SetattrNested(request, current_token_attribute, None) 107 while limit is None or limit: 108 if batch_size_attribute: 109 # On Py3, None is not comparable so min() below will fail. 110 # On Py2, None is always less than any number so if batch_size 111 # is None, the request_batch_size will always be None regardless 112 # of the value of limit. This doesn't generally strike me as the 113 # correct behavior, but this change preserves the existing Py2 114 # behavior on Py3. 115 if batch_size is None: 116 request_batch_size = None 117 else: 118 request_batch_size = min(batch_size, limit or batch_size) 119 _SetattrNested(request, batch_size_attribute, request_batch_size) 120 response = getattr(service, method)(request, 121 global_params=global_params) 122 items = get_field_func(response, field) 123 if predicate: 124 items = list(filter(predicate, items)) 125 for item in items: 126 yield item 127 if limit is None: 128 continue 129 limit -= 1 130 if not limit: 131 return 132 token = _GetattrNested(response, next_token_attribute) 133 if not token: 134 return 135 _SetattrNested(request, current_token_attribute, token) 136