Move the page-related call() methods and add tests (#1053)

* Move callGAPIItems and add unit tests

* Move the page-related call() methods and add tests

Continues work on jay0lee/GAM#147

* Account for rand addition that rounds to 1.0

In this case, sleep time could be equal to 61

src/gapi/__init__.py

@@ -1,5 +1,7 @@
 """Methods related to execution of GAPI requests."""
 
+import sys
+
 import controlflow
 import display
 from gapi import errors
@@ -8,8 +10,8 @@ import httplib2
 from var import (GC_CA_FILE, GC_Values, GC_TLS_MIN_VERSION, GC_TLS_MAX_VERSION,
                  GM_Globals, GM_CURRENT_API_SCOPES, GM_CURRENT_API_USER,
                  GM_EXTRA_ARGS_DICT, GM_OAUTH2SERVICE_ACCOUNT_CLIENT_ID,
-                 MESSAGE_API_ACCESS_CONFIG, MESSAGE_API_ACCESS_DENIED,
-                 MESSAGE_SERVICE_NOT_APPLICABLE)
+                 MAX_RESULTS_API_EXCEPTIONS, MESSAGE_API_ACCESS_CONFIG,
+                 MESSAGE_API_ACCESS_DENIED, MESSAGE_SERVICE_NOT_APPLICABLE)
 import google.auth.exceptions
 
 
@@ -138,6 +140,163 @@ def call(service,
       controlflow.system_error_exit(4, str(e))
 
 
+def get_items(service,
+              function,
+              items='items',
+              throw_reasons=None,
+              retry_reasons=None,
+              **kwargs):
+  """Gets a single page of items from a Google service function that is paged.
+
+  Args:
+    service: A Google service object for the desired API.
+    function: String, The name of a service request method to execute.
+    items: String, the name of the resulting "items" field within the service
+      method's response object.
+    throw_reasons: A list of Google HTTP error reason strings indicating the
+      errors generated by this request should be re-thrown. All other HTTP
+      errors are consumed.
+    retry_reasons: A list of Google HTTP error reason strings indicating which
+      error should be retried, using exponential backoff techniques, when the
+      error reason is encountered.
+    **kwargs: Additional params to pass to the request method.
+
+  Returns:
+    The list of items in the first page of a response.
+  """
+  results = call(
+      service,
+      function,
+      throw_reasons=throw_reasons,
+      retry_reasons=retry_reasons,
+      **kwargs)
+  if results:
+    return results.get(items, [])
+  return []
+
+
+def _get_max_page_size_for_api_call(service, function, **kwargs):
+  """Gets the maximum number of results supported for a single API call.
+
+  Args:
+    service: A Google service object for the desired API.
+    function: String, The name of the service method to check for max page size.
+    **kwargs: Additional params that will be passed to the request method.
+
+  Returns:
+    Int, A value from discovery if it exists, otherwise value from
+        MAX_RESULTS_API_EXCEPTIONS, otherwise None
+  """
+  method = getattr(service, function)
+  api_id = method(**kwargs).methodId
+  for resource in service._rootDesc.get('resources', {}).values():
+    for a_method in resource.get('methods', {}).values():
+      if a_method.get('id') == api_id:
+        if not a_method.get('parameters') or a_method['parameters'].get(
+            'pageSize') or not a_method['parameters'].get('maxResults'):
+          # Make sure API call supports maxResults. For now we don't care to
+          # set pageSize since all known pageSize API calls have
+          # default pageSize == max pageSize.
+          return None
+        known_api_max = MAX_RESULTS_API_EXCEPTIONS.get(api_id)
+        max_results = a_method['parameters']['maxResults'].get(
+            'maximum', known_api_max)
+        return {'maxResults': max_results}
+
+  return None
+
+
+def get_all_pages(service,
+                  function,
+                  items='items',
+                  page_message=None,
+                  message_attribute=None,
+                  soft_errors=False,
+                  throw_reasons=None,
+                  retry_reasons=None,
+                  **kwargs):
+  """Aggregates and returns all pages of a Google service function response.
+
+  All pages of items are aggregated and returned as a single list.
+
+  Args:
+    service: A Google service object for the desired API.
+    function: String, The name of a service request method to execute.
+    items: String, the name of the resulting "items" field within the method's
+      response object. The items in this field will be aggregated across all
+      pages and returned.
+    page_message: String, a message to be displayed to the user during paging.
+      Template strings allow for dynamic content to be inserted during paging.
+        Supported template strings:
+          %%total_items%% : The current number of items discovered across all
+            pages.
+          %%first_item%%  : In conjunction with `message_attribute` arg, will
+            display a unique property of the first item in the current page.
+          %%last_item%%   : In conjunction with `message_attribute` arg, will
+            display a unique property of the last item in the current page.
+    message_attribute: String, the name of a signature field within a single
+      returned item which identifies that unique item. This field is used with
+      `page_message` to templatize a paging status message.
+    soft_errors: Bool, If True, writes non-fatal errors to stderr.
+    throw_reasons: A list of Google HTTP error reason strings indicating the
+      errors generated by this request should be re-thrown. All other HTTP
+      errors are consumed.
+    retry_reasons: A list of Google HTTP error reason strings indicating which
+      error should be retried, using exponential backoff techniques, when the
+      error reason is encountered.
+    **kwargs: Additional params to pass to the request method.
+
+  Returns:
+    A list of all items received from all paged responses.
+  """
+  if 'maxResults' not in kwargs and 'pageSize' not in kwargs:
+    page_key = _get_max_page_size_for_api_call(service, function, **kwargs)
+    if page_key:
+      kwargs.update(page_key)
+  all_items = []
+  page_token = None
+  total_items = 0
+  while True:
+    page = call(
+        service,
+        function,
+        soft_errors=soft_errors,
+        throw_reasons=throw_reasons,
+        retry_reasons=retry_reasons,
+        pageToken=page_token,
+        **kwargs)
+    if page:
+      page_token = page.get('nextPageToken')
+      page_items = page.get(items, [])
+      num_page_items = len(page_items)
+      total_items += num_page_items
+      all_items.extend(page_items)
+    else:
+      page_token = None
+      num_page_items = 0
+
+    # Show a paging message to the user that indicates paging progress
+    if page_message:
+      show_message = page_message.replace('%%total_items%%', str(total_items))
+      if message_attribute:
+        first_item = page_items[0] if num_page_items > 0 else {}
+        last_item = page_items[-1] if num_page_items > 1 else first_item
+        show_message = show_message.replace(
+            '%%first_item%%', str(first_item.get(message_attribute, '')))
+        show_message = show_message.replace(
+            '%%last_item%%', str(last_item.get(message_attribute, '')))
+      sys.stderr.write('\r')
+      sys.stderr.flush()
+      sys.stderr.write(show_message)
+
+    if not page_token:
+      # End the paging status message and return all items.
+      if page_message and (page_message[-1] != '\n'):
+        sys.stderr.write('\r\n')
+        sys.stderr.flush()
+      return all_items
+
+
 # TODO: Make this private once all execution related items that use this method
 # have been brought into this file
 def handle_oauth_token_error(e, soft_errors):