1# -*- coding: utf-8 -*- #
2# Copyright 2014 Google LLC. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#    http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15""""Helpers for making batch requests."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import unicode_literals
20
21import json
22
23from apitools.base.py import batch
24from apitools.base.py import exceptions
25
26from googlecloudsdk.api_lib.util import apis
27from googlecloudsdk.core import properties
28
29# Upper bound on batch size
30# https://cloud.google.com/compute/docs/api/how-tos/batch
31_BATCH_SIZE_LIMIT = 1000
32
33
34class BatchChecker(object):
35  """Class to conveniently curry the prompted_service_tokens cache."""
36
37  def __init__(self, prompted_service_tokens):
38    """Initialize class.
39
40    Args:
41      prompted_service_tokens: a set of string tokens that have already been
42        prompted for enablement.
43    """
44    self.prompted_service_tokens = prompted_service_tokens
45
46  # pylint: disable=unused-argument
47  def BatchCheck(self, http_response, exception):
48    """Callback for apitools batch responses.
49
50    This will use self.prompted_service_tokens to cache service tokens that
51    have already been prompted. In this way, if the same service has multiple
52    batch requests and is enabled on the first, the user won't get a bunch of
53    superflous messages. Note that this cannot be reused between batch uses
54    because of the mutation.
55
56    Args:
57      http_response: Deserialized http_wrapper.Response object.
58      exception: apiclient.errors.HttpError object if an error occurred.
59    """
60    # If there is no exception, then there is not an api enablement error.
61    # Also, if prompting to enable is disabled, then we let the batch module
62    # fail the batch request.
63    if (exception is None
64        or not properties.VALUES.core.should_prompt_to_enable_api.GetBool()):
65      return
66    enablement_info = apis.GetApiEnablementInfo(exception)
67    if not enablement_info:  # Exception was not an api enablement error.
68      return
69    project, service_token, exception = enablement_info
70    if service_token not in self.prompted_service_tokens:  # Only prompt once.
71      self.prompted_service_tokens.add(service_token)
72      apis.PromptToEnableApi(project, service_token, exception,
73                             is_batch_request=True)
74
75
76def MakeRequests(requests, http, batch_url=None):
77  """Makes batch requests.
78
79  Args:
80    requests: A list of tuples. Each tuple must be of the form
81        (service, method, request object).
82    http: An HTTP object.
83    batch_url: The URL to which to send the requests.
84
85  Returns:
86    A tuple where the first element is a list of all objects returned
87    from the calls and the second is a list of error messages.
88  """
89  retryable_codes = []
90  if properties.VALUES.core.should_prompt_to_enable_api.GetBool():
91    # If the compute API is not enabled, then a 403 error is returned. We let
92    # the batch module handle retrying requests by adding 403 to the list of
93    # retryable codes for the batch request. If we should not prompt, then
94    # we keep retryable_codes empty, so the request fails.
95    retryable_codes.append(apis.API_ENABLEMENT_ERROR_EXPECTED_STATUS_CODE)
96  batch_request = batch.BatchApiRequest(batch_url=batch_url,
97                                        retryable_codes=retryable_codes)
98  for service, method, request in requests:
99    batch_request.Add(service, method, request)
100
101  # TODO(b/36030477) this shouldn't be necessary in the future when batch and
102  # non-batch error handling callbacks are unified
103  batch_checker = BatchChecker(set())
104  responses = batch_request.Execute(
105      http, max_batch_size=_BATCH_SIZE_LIMIT,
106      batch_request_callback=batch_checker.BatchCheck)
107
108  objects = []
109  errors = []
110
111  for response in responses:
112    objects.append(response.response)
113
114    if response.is_error:
115      # TODO(b/33771874): Use HttpException to decode error payloads.
116      error_message = None
117      if isinstance(response.exception, exceptions.HttpError):
118        try:
119          data = json.loads(response.exception.content)
120          error_message = (
121              response.exception.status_code,
122              data.get('error', {}).get('message'))
123        except ValueError:
124          pass
125        if not error_message:
126          error_message = (response.exception.status_code,
127                           response.exception.content)
128      else:
129        error_message = (None, response.exception.message)
130
131      errors.append(error_message)
132
133  return objects, errors
134