| # Copyright 2014 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. |
| |
| """Library for gcloud operations.""" |
| |
| from __future__ import print_function |
| |
| import logging |
| |
| from chromite.lib import cros_build_lib |
| from chromite.lib import timeout_util |
| |
| |
| class GCContextException(Exception): |
| """Base exception for this module.""" |
| |
| |
| class ZoneNotSpecifiedError(GCContextException): |
| """Raised when zone is not specified for a zone-specific command.""" |
| |
| |
| class GCContext(object): |
| """A wrapper around the gcloud commandline tool. |
| |
| Currently supports only `gcloud compute`. |
| """ |
| GCLOUD_BASE_COMMAND = 'gcloud' |
| GCLOUD_COMPUTE_COMMAND = 'compute' |
| |
| def __init__(self, project, zone=None, quiet=False): |
| """Initializes GCContext. |
| |
| Args: |
| project: The Google Cloud project to use. |
| zone: The default zone to operate on when zone is not given for a |
| zone-specific command. |
| quiet: If set True, skip any user prompts and use the default value. |
| """ |
| self.project = project |
| self.zone = zone |
| self.quiet = quiet |
| |
| @classmethod |
| def _GetBaseComputeCommand(cls): |
| """Returns the base Google Compute Engine command.""" |
| return [cls.GCLOUD_BASE_COMMAND, cls.GCLOUD_COMPUTE_COMMAND] |
| |
| def DoCommand(self, cmd, **kwargs): |
| """Runs |cmd|. |
| |
| cmd: The command to run. |
| project: The project to use. Defaults to using self.project. |
| kwargs: See cros_build_lib.RunCommand. |
| """ |
| cmd = self._GetBaseComputeCommand() + cmd |
| cmd += ['--project', kwargs.pop('project', self.project)] |
| if kwargs.pop('quiet', self.quiet): |
| cmd += ['--quiet'] |
| zone = kwargs.pop('zone', None) |
| if zone: |
| cmd += ['--zone', zone] |
| |
| return cros_build_lib.RunCommand(cmd, **kwargs) |
| |
| def DoZoneSpecificCommand(self, cmd, **kwargs): |
| """Runs the zone-specific |cmd|. |
| |
| cmd: The command to run. |
| zone: The zone to use. Defaults to using self.zone. |
| kwargs: See DoCommand. |
| """ |
| kwargs.setdefault('zone', self.zone) |
| if not kwargs.get('zone'): |
| # To avoid ambiguity, force user to specify a zone (or a default |
| # zone) when accessing zone-specific resources. |
| raise ZoneNotSpecifiedError() |
| return self.DoCommand(cmd, **kwargs) |
| |
| def CopyFilesToInstance(self, instance, src, dest, **kwargs): |
| """Copies files from |src| to |dest| on |instance|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path. |
| dest: The destination path. |
| kwargs: See DoCommand. |
| """ |
| return self._CopyFiles(src, '%s:%s' % (instance, dest), **kwargs) |
| |
| def CopyFilesFromInstance(self, instance, src, dest, **kwargs): |
| """Copies files from |src| on |instance| to local |dest|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path. |
| dest: The destination path. |
| kwargs: See DoCommand. |
| """ |
| return self._CopyFiles('%s:%s' % (instance, src), dest, **kwargs) |
| |
| def _CopyFiles(self, src, dest, **kwargs): |
| """Copies files from |src| to |dest|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path (a local path or instance@path). |
| dest: The destination path (a local path or instance@path). |
| kwargs: See DoCommand. |
| """ |
| return self.DoZoneSpecificCommand(['copy-files', src, dest], **kwargs) |
| |
| def SSH(self, instance, cmd=None, **kwargs): |
| """SSH into |instance|. Run |cmd| if it is provided. |
| |
| Args: |
| instance: Name of the instance. |
| cmd: Command to run on |instance|. |
| """ |
| ssh_cmd = ['ssh', instance] |
| if cmd: |
| ssh_cmd += ['--command', cmd] |
| return self.DoZoneSpecificCommand(ssh_cmd, **kwargs) |
| |
| def ListInstances(self, **kwargs): |
| """Lists all instances.""" |
| return self.DoCommand(['instances', 'list'], **kwargs) |
| |
| def ListDisks(self, **kwargs): |
| """Lists all disks.""" |
| return self.DoCommand(['disks', 'list'], **kwargs) |
| |
| def ListImages(self, **kwargs): |
| """Lists all instances.""" |
| return self.DoCommand(['images', 'list'], **kwargs) |
| |
| def CreateImage(self, image, source_uri=None, disk=None, **kwargs): |
| """Creates an image from |source_uri| or |disk|. |
| |
| Args: |
| image: The name of image to create. |
| source_uri: The tar.gz image file (e.g. gs://foo/bar/image.tar.gz) |
| disk: The source disk to create the image from. One and only one of |
| |source_uril| and |disk| should be set. |
| kwargs: See DoCommand. |
| """ |
| if source_uri and disk: |
| raise GCContextException('Cannot specify both source uri and disk.') |
| |
| cmd = ['images', 'create', image] |
| if disk: |
| cmd += ['--source-disk', disk] |
| zone = kwargs.get('zone', self.zone) |
| if zone: |
| # Disks are zone-specific resources. |
| cmd += ['--source-disk-zone', zone] |
| |
| if source_uri: |
| cmd += ['--source-uri', source_uri] |
| |
| return self.DoCommand(cmd, **kwargs) |
| |
| def CreateInstance(self, instance, image=None, machine_type=None, |
| address=None, wait_until_sshable=True, |
| scopes=None, **kwargs): |
| """Creates an |instance|. |
| |
| Args: |
| instance: The name of the instance to create. |
| image: The source image to create |instance| from. |
| machine_type: The machine type to use. |
| address: The external IP address to assign to |instance|. |
| wait_until_sshable: After creating |instance|, wait until |
| we can ssh into |instance|. |
| scopes: The list (or tuple) of service account scopes. |
| kwargs: See DoZoneSpecificCommand. |
| """ |
| cmd = ['instances', 'create', instance] |
| if image: |
| cmd += ['--image', image] |
| if address: |
| cmd += ['--address', address] |
| if machine_type: |
| cmd += ['--machine-type', machine_type] |
| if scopes: |
| cmd = cmd + ['--scopes'] + list(scopes) |
| |
| ret = self.DoZoneSpecificCommand(cmd, **kwargs) |
| if wait_until_sshable: |
| def _IsUp(): |
| try: |
| self.SSH(instance, cmd='ls', capture_output=True) |
| except cros_build_lib.RunCommandError: |
| return False |
| else: |
| return True |
| |
| try: |
| logging.info('Waiting for the instance to be sshable...') |
| timeout = 60 * 5 |
| timeout_util.WaitForReturnTrue(_IsUp, timeout, period=5) |
| except timeout_util.TimeoutError: |
| raise GCContextException('Timed out wating to ssh into the instance') |
| |
| return ret |
| |
| def DeleteInstance(self, instance, keep_disks=None, **kwargs): |
| """Deletes |instance|. User will be prompted to confirm. |
| |
| Args: |
| instance: Name of the instance. |
| keep_disks: Keep the type of the disk; valid types are |
| 'boot', 'data', and 'all'. |
| kwargs: See DoCommand. |
| """ |
| cmd = ['instances', 'delete', instance] |
| if keep_disks: |
| cmd += ['--keep-disks', keep_disks] |
| |
| return self.DoZoneSpecificCommand(cmd, **kwargs) |
| |
| def DeleteImage(self, image, **kwargs): |
| """Deletes |image|. User will be prompted to confirm. |
| |
| Args: |
| image: Name of the image. |
| kwargs: See DoCommand. |
| """ |
| return self.DoCommand(['images', 'delete', image], **kwargs) |
| |
| def DeleteDisk(self, disk, **kwargs): |
| """Deletes |disk|. |
| |
| Args: |
| disk: Name of the disk. |
| kwargs: See DoCommand. |
| """ |
| cmd = ['disks', 'delete', disk] |
| return self.DoZoneSpecificCommand(cmd, **kwargs) |
| |