| # Copyright 2015 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """A convinient wrapper of the GCE python API. |
| |
| Public methods in class GceContext raise HttpError when the underlining call to |
| Google API fails, or gce.Error on other failures. |
| """ |
| |
| from __future__ import print_function |
| |
| import httplib2 |
| |
| from chromite.lib import cros_logging as logging |
| from chromite.lib import timeout_util |
| from googleapiclient.discovery import build |
| from googleapiclient.errors import HttpError |
| from googleapiclient.http import HttpRequest |
| from oauth2client.client import GoogleCredentials |
| |
| |
| class Error(Exception): |
| """Base exception for this module.""" |
| |
| |
| class ResourceNotFoundError(Error): |
| """Exceptions raised when requested GCE resource was not found.""" |
| |
| |
| class RetryOnServerErrorHttpRequest(HttpRequest): |
| """A HttpRequest that will be retried on server errors automatically.""" |
| |
| def __init__(self, num_retries, *args, **kwargs): |
| """Constructor for RetryOnServerErrorHttpRequest.""" |
| self.num_retries = num_retries |
| super(RetryOnServerErrorHttpRequest, self).__init__(*args, **kwargs) |
| |
| def execute(self, http=None, num_retries=None): |
| """Excutes a RetryOnServerErrorHttpRequest. |
| |
| HttpRequest.execute() has the option of automatically retrying on server |
| errors, i.e., 500 status codes. Call it with a non-zero value of |
| |num_retries| will cause failed requests to be retried. |
| |
| Args: |
| http: The httplib2.http to send this request through. |
| num_retries: Number of retries. Class default value will be used if |
| omitted. |
| |
| Returns: |
| A deserialized object model of the response body as determined |
| by the postproc. See HttpRequest.execute(). |
| """ |
| return super(RetryOnServerErrorHttpRequest, self).execute( |
| http=http, num_retries=num_retries or self.num_retries) |
| |
| |
| class GceContext(object): |
| """A convinient wrapper around the GCE Python API.""" |
| |
| _GCE_SCOPES = ( |
| 'https://www.googleapis.com/auth/compute', # CreateInstance, CreateImage |
| 'https://www.googleapis.com/auth/devstorage.full_control', # CreateImage |
| ) |
| _DEFAULT_NETWORK = 'default' |
| _DEFAULT_MACHINE_TYPE = 'n1-standard-8' |
| _DEFAULT_TIMEOUT_SEC = 5 * 60 |
| _INSTANCE_OPERATIONS_TIMEOUT_SEC = 5 * 60 |
| _IMAGE_OPERATIONS_TIMEOUT_SEC = 2* 60 |
| RETRIES = 2 |
| |
| def __init__(self, project, zone, credentials, thread_safe=False): |
| """Initializes GceContext. |
| |
| Args: |
| project: The GCP project to create instances in. |
| zone: The default zone to create instances in. |
| credentials: The credentials used to call the GCE API. |
| thread_safe: Whether the client is expected to be thread safe. |
| """ |
| self.project = project |
| self.zone = zone |
| |
| def _BuildRequest(http, *args, **kwargs): |
| """Custom request builder.""" |
| return self._BuildRetriableRequest(self.RETRIES, http, thread_safe, |
| credentials, *args, **kwargs) |
| |
| self.gce_client = build('compute', 'v1', credentials=credentials, |
| requestBuilder=_BuildRequest) |
| |
| @classmethod |
| def ForServiceAccount(cls, project, zone, json_key_file): |
| """Creates a GceContext using service account credentials. |
| |
| About service account: |
| https://developers.google.com/api-client-library/python/auth/service-accounts |
| |
| Args: |
| project: The GCP project to create images and instances in. |
| zone: The default zone to create instances in. |
| json_key_file: Path to the service account JSON key. |
| |
| Returns: |
| GceContext. |
| """ |
| credentials = GoogleCredentials.from_stream(json_key_file).create_scoped( |
| cls._GCE_SCOPES) |
| return GceContext(project, zone, credentials) |
| |
| @classmethod |
| def ForServiceAccountThreadSafe(cls, project, zone, json_key_file): |
| """Creates a thread-safe GceContext using service account credentials. |
| |
| About service account: |
| https://developers.google.com/api-client-library/python/auth/service-accounts |
| |
| Args: |
| project: The GCP project to create images and instances in. |
| zone: The default zone to create instances in. |
| json_key_file: Path to the service account JSON key. |
| |
| Returns: |
| GceContext. |
| """ |
| credentials = GoogleCredentials.from_stream(json_key_file).create_scoped( |
| cls._GCE_SCOPES) |
| return GceContext(project, zone, credentials, thread_safe=True) |
| |
| def CreateInstance(self, name, image, zone=None, network=None, |
| machine_type=None, **kwargs): |
| """Creates an instance with the given image and waits until it's ready. |
| |
| Args: |
| name: Instance name. |
| image: Fully spelled URL of the image, e.g., for private images, |
| 'global/images/my-private-image', or for images from a |
| publicly-available project, |
| 'projects/debian-cloud/global/images/debian-7-wheezy-vYYYYMMDD'. |
| Details: |
| https://cloud.google.com/compute/docs/reference/latest/instances/insert |
| zone: The zone to create the instance in. Default zone will be used if |
| omitted. |
| network: An existing network to create the instance in. Default network |
| will be used if omitted. |
| machine_type: The machine type to use. Default machine type will be used |
| if omitted. |
| kwargs: Other possible Instance Resource properties. |
| https://cloud.google.com/compute/docs/reference/latest/instances#resource |
| |
| Returns: |
| URL to the created instance. |
| """ |
| machine_type = 'zones/%s/machineTypes/%s' % ( |
| zone or self.zone, machine_type or self._DEFAULT_MACHINE_TYPE) |
| # Allow machineType overriding. |
| if 'machineType' in kwargs: |
| machine_type = kwargs['machineType'] |
| |
| config = { |
| 'name': name, |
| 'machineType': machine_type, |
| 'disks': ( |
| { |
| 'boot': True, |
| 'autoDelete': True, |
| 'initializeParams': { |
| 'sourceImage': image, |
| }, |
| }, |
| ), |
| 'networkInterfaces': ( |
| { |
| 'network': ('global/networks/%s' % network or |
| self._DEFAULT_NETWORK), |
| 'accessConfigs': ( |
| { |
| 'type': 'ONE_TO_ONE_NAT', |
| 'name': 'External NAT', |
| }, |
| ), |
| }, |
| ), |
| } |
| config.update(**kwargs) |
| operation = self.gce_client.instances().insert( |
| project=self.project, |
| zone=zone or self.zone, |
| body=config).execute() |
| self._WaitForZoneOperation( |
| operation['name'], |
| timeout_sec=self._INSTANCE_OPERATIONS_TIMEOUT_SEC, |
| timeout_handler=lambda: self.DeleteInstance(name)) |
| return operation['targetLink'] |
| |
| def DeleteInstance(self, name, zone=None): |
| """Deletes an instance with the name and waits until it's done. |
| |
| Args: |
| name: Name of the instance to delete. |
| zone: Zone where the instance is in. Default zone will be used if omitted. |
| """ |
| operation = self.gce_client.instances().delete( |
| project=self.project, |
| zone=zone or self.zone, |
| instance=name).execute() |
| self._WaitForZoneOperation( |
| operation['name'], timeout_sec=self._INSTANCE_OPERATIONS_TIMEOUT_SEC) |
| |
| def CreateImage(self, name, source): |
| """Creates an image with the given |source|. |
| |
| Args: |
| name: Name of the image to be created. |
| source: |
| Google Cloud Storage object of the source disk, e.g., |
| 'https://storage.googleapis.com/my-gcs-bucket/test_image.tar.gz'. |
| |
| Returns: |
| URL to the created image. |
| """ |
| config = { |
| 'name': name, |
| 'rawDisk': { |
| 'source': source, |
| }, |
| } |
| operation = self.gce_client.images().insert( |
| project=self.project, |
| body=config).execute() |
| self._WaitForGlobalOperation(operation['name'], |
| timeout_sec=self._IMAGE_OPERATIONS_TIMEOUT_SEC, |
| timeout_handler=lambda: self.DeleteImage(name)) |
| return operation['targetLink'] |
| |
| def DeleteImage(self, name): |
| """Deletes an image and waits until it's deleted. |
| |
| Args: |
| name: Name of the image to delete. |
| """ |
| operation = self.gce_client.images().delete( |
| project=self.project, |
| image=name).execute() |
| self._WaitForGlobalOperation(operation['name'], |
| timeout_sec=self._IMAGE_OPERATIONS_TIMEOUT_SEC) |
| |
| def ListInstances(self, zone=None): |
| """Lists all instances. |
| |
| Args: |
| zone: Zone where the instances are in. Default zone will be used if |
| omitted. |
| |
| Returns: |
| A list of Instance Resources if found, or an empty list otherwise. |
| """ |
| result = self.gce_client.instances().list(project=self.project, |
| zone=zone or self.zone).execute() |
| return result.get('items', []) |
| |
| def ListImages(self): |
| """Lists all images. |
| |
| Returns: |
| A list of Image Resources if found, or an empty list otherwise. |
| """ |
| result = self.gce_client.images().list(project=self.project).execute() |
| return result.get('items', []) |
| |
| def GetInstance(self, instance, zone=None): |
| """Gets an Instance Resource by name and zone. |
| |
| Args: |
| instance: Name of the instance. |
| zone: Zone where the instance is in. Default zone will be used if omitted. |
| |
| Returns: |
| An Instance Resource. |
| |
| Raises: |
| ResourceNotFoundError if instance was not found, or HttpError on other |
| HTTP failures. |
| """ |
| try: |
| return self.gce_client.instances().get(project=self.project, |
| zone=zone or self.zone, |
| instance=instance).execute() |
| except HttpError as e: |
| if e.resp.status == 404: |
| raise ResourceNotFoundError( |
| 'Instance "%s" for project "%s" in zone "%s" was not found.' % |
| (instance, self.project, zone or self.zone)) |
| else: |
| raise |
| |
| def GetInstanceIP(self, instance, zone=None): |
| """Gets the external IP of an instance. |
| |
| Args: |
| instance: Name of the instance to get IP for. |
| zone: Zone where the instance is in. Default zone will be used if omitted. |
| |
| Returns: |
| External IP address of the instance. |
| |
| Raises: |
| Error: Something went wrong when trying to get IP for the instance. |
| """ |
| result = self.GetInstance(instance, zone) |
| try: |
| return result['networkInterfaces'][0]['accessConfigs'][0]['natIP'] |
| except (KeyError, IndexError): |
| raise Error('Failed to get IP address for instance %s' % instance) |
| |
| def GetImage(self, image): |
| """Gets an Image Resource by name. |
| |
| Args: |
| image: Name of the image to look for. |
| |
| Returns: |
| An Image Resource. |
| |
| Raises: |
| ResourceNotFoundError: The requested image was not found. |
| """ |
| try: |
| return self.gce_client.images().get(project=self.project, |
| image=image).execute() |
| except HttpError as e: |
| if e.resp.status == 404: |
| raise ResourceNotFoundError('Image "%s" for project "%s" was not found.' |
| % (image, self.project)) |
| else: |
| raise |
| |
| def InstanceExists(self, instance, zone=None): |
| """Checks if an instance exists in the current project. |
| |
| Args: |
| instance: Name of the instance to check existence of. |
| zone: Zone where the instance is in. Default zone will be used if omitted. |
| |
| Returns: |
| True if the instance exists or False otherwise. |
| """ |
| try: |
| return self.GetInstance(instance, zone) is not None |
| except ResourceNotFoundError: |
| return False |
| |
| def ImageExists(self, image): |
| """Checks if an image exists in the current project. |
| |
| Args: |
| image: Name of the image to check existence of. |
| |
| Returns: |
| True if the instance exists or False otherwise. |
| """ |
| try: |
| return self.GetImage(image) is not None |
| except ResourceNotFoundError: |
| return False |
| |
| def _WaitForZoneOperation(self, operation, zone=None, timeout_sec=None, |
| timeout_handler=None): |
| """Waits until a GCE ZoneOperation is finished or timed out. |
| |
| Args: |
| operation: The GCE operation to wait for. |
| zone: The zone that |operation| belongs to. |
| timeout_sec: The maximum number of seconds to wait for. |
| timeout_handler: A callable to be executed when timeout happens. |
| """ |
| get_request = self.gce_client.zoneOperations().get( |
| project=self.project, zone=zone or self.zone, operation=operation) |
| self._WaitForOperation(operation, get_request, timeout_sec, |
| timeout_handler=timeout_handler) |
| |
| def _WaitForGlobalOperation(self, operation, timeout_sec=None, |
| timeout_handler=None): |
| """Waits until a GCE GlobalOperation is finished or timed out. |
| |
| Args: |
| operation: The GCE operation to wait for. |
| timeout_sec: The maximum number of seconds to wait for. |
| timeout_handler: A callable to be executed when timeout happens. |
| """ |
| get_request = self.gce_client.globalOperations().get(project=self.project, |
| operation=operation) |
| self._WaitForOperation(operation, get_request, timeout_sec=timeout_sec, |
| timeout_handler=timeout_handler) |
| |
| def _WaitForOperation(self, operation, get_operation_request, |
| timeout_sec=None, timeout_handler=None): |
| """Waits until timeout or the request gets a response with a 'DONE' status. |
| |
| Args: |
| operation: The GCE operation to wait for. |
| get_operation_request: |
| The HTTP request to get the operation's status. |
| This request will be executed periodically until it returns a status |
| 'DONE'. |
| timeout_sec: The maximum number of seconds to wait for. |
| timeout_handler: A callable to be executed when times out. |
| """ |
| def _IsDone(): |
| result = get_operation_request.execute() |
| if result['status'] == 'DONE': |
| if 'error' in result: |
| raise Error(result['error']) |
| return True |
| return False |
| |
| try: |
| timeout = timeout_sec or self._DEFAULT_TIMEOUT_SEC |
| logging.info('Waiting up to %d seconds for operation [%s] to complete...', |
| timeout, operation) |
| timeout_util.WaitForReturnTrue(_IsDone, timeout, period=1) |
| except timeout_util.TimeoutError: |
| if not timeout_handler: |
| timeout_handler() |
| raise Error('Timeout wating for operation [%s] to complete' % operation) |
| |
| def _BuildRetriableRequest(self, num_retries, http, thread_safe=False, |
| credentials=None, *args, **kwargs): |
| """Builds a request that will be automatically retried on server errors. |
| |
| Args: |
| num_retries: The maximum number of times to retry until give up. |
| http: An httplib2.Http object that this request will be executed through. |
| thread_safe: Whether or not the request needs to be thread-safe. |
| credentials: Credentials to apply to the request. |
| *args: Optional positional arguments. |
| **kwargs: Optional keyword arguments. |
| |
| Returns: |
| RetryOnServerErrorHttpRequest: A request that will automatically retried |
| on server errors. |
| """ |
| if thread_safe: |
| # Create a new http object for every request. |
| http = credentials.authorize(httplib2.Http()) |
| return RetryOnServerErrorHttpRequest(num_retries, http, *args, **kwargs) |
