~andrewjbeach/juju-ci-tools/make-local-patcher

from __future__ import print_function

from collections import (

    defaultdict,

    namedtuple,

)

from contextlib import (

    contextmanager,

    nested,

)

from copy import deepcopy

from cStringIO import StringIO

from datetime import timedelta

import errno

from itertools import chain

import logging

import os

import re

from shutil import rmtree

import subprocess

import sys

from tempfile import NamedTemporaryFile

import time

import yaml

from jujuconfig import (

    get_environments_path,

    get_jenv_path,

    get_juju_home,

    get_selected_environment,

)

from utility import (

    check_free_disk_space,

    ensure_deleted,

    ensure_dir,

    is_ipv6_address,

    pause,

    scoped_environ,

    split_address_port,

    temp_dir,

    until_timeout,

)

__metaclass__ = type

AGENTS_READY = set(['started', 'idle'])

WIN_JUJU_CMD = os.path.join('\\', 'Progra~2', 'Juju', 'juju.exe')

JUJU_DEV_FEATURE_FLAGS = 'JUJU_DEV_FEATURE_FLAGS'

CONTROLLER = 'controller'

KILL_CONTROLLER = 'kill-controller'

SYSTEM = 'system'

_DEFAULT_BUNDLE_TIMEOUT = 3600

_jes_cmds = {KILL_CONTROLLER: {

    'create': 'create-environment',

    'kill': KILL_CONTROLLER,

}}

for super_cmd in [SYSTEM, CONTROLLER]:

    _jes_cmds[super_cmd] = {

        'create': '{} create-environment'.format(super_cmd),

        'kill': '{} kill'.format(super_cmd),

    }

log = logging.getLogger("jujupy")

def get_timeout_path():

    import timeout

    return os.path.abspath(timeout.__file__)

def get_timeout_prefix(duration, timeout_path=None):

    """Return extra arguments to run a command with a timeout."""

    if timeout_path is None:

        timeout_path = get_timeout_path()

    return (sys.executable, timeout_path, '%.2f' % duration, '--')

def parse_new_state_server_from_error(error):

    err_str = str(error)

    output = getattr(error, 'output', None)

    if output is not None:

        err_str += output

    matches = re.findall(r'Attempting to connect to (.*):22', err_str)

    if matches:

        return matches[-1]

    return None

class ErroredUnit(Exception):

    def __init__(self, unit_name, state):

        msg = '%s is in state %s' % (unit_name, state)

        Exception.__init__(self, msg)

        self.unit_name = unit_name

        self.state = state

class BootstrapMismatch(Exception):

    def __init__(self, arg_name, arg_val, env_name, env_val):

        super(BootstrapMismatch, self).__init__(

            '--{} {} does not match {}: {}'.format(

                arg_name, arg_val, env_name, env_val))

class UpgradeMongoNotSupported(Exception):

    def __init__(self):

        super(UpgradeMongoNotSupported, self).__init__(

            'This client does not support upgrade-mongo')

class JESNotSupported(Exception):

    def __init__(self):

        super(JESNotSupported, self).__init__(

            'This client does not support JES')

class JESByDefault(Exception):

    def __init__(self):

        super(JESByDefault, self).__init__(

            'This client does not need to enable JES')

Machine = namedtuple('Machine', ['machine_id', 'info'])

def yaml_loads(yaml_str):

    return yaml.safe_load(StringIO(yaml_str))

def coalesce_agent_status(agent_item):

    """Return the machine agent-state or the unit agent-status."""

    state = agent_item.get('agent-state')

    if state is None and agent_item.get('agent-status') is not None:

        state = agent_item.get('agent-status').get('current')

    if state is None and agent_item.get('juju-status') is not None:

        state = agent_item.get('juju-status').get('current')

    if state is None:

        state = 'no-agent'

    return state

def make_client(juju_path, debug, env_name, temp_env_name):

    env = SimpleEnvironment.from_config(env_name)

    if temp_env_name is not None:

        env.set_model_name(temp_env_name)

    return EnvJujuClient.by_version(env, juju_path, debug)

class CannotConnectEnv(subprocess.CalledProcessError):

    def __init__(self, e):

        super(CannotConnectEnv, self).__init__(e.returncode, e.cmd, e.output)

class StatusNotMet(Exception):

    _fmt = 'Expected status not reached in {env}.'

    def __init__(self, environment_name, status):

        self.env = environment_name

        self.status = status

    def __str__(self):

        return self._fmt.format(env=self.env)

class AgentsNotStarted(StatusNotMet):

    _fmt = 'Timed out waiting for agents to start in {env}.'

class VersionsNotUpdated(StatusNotMet):

    _fmt = 'Some versions did not update.'

class WorkloadsNotReady(StatusNotMet):

    _fmt = 'Workloads not ready in {env}.'

@contextmanager

def temp_yaml_file(yaml_dict):

    temp_file = NamedTemporaryFile(suffix='.yaml', delete=False)

    try:

        with temp_file:

            yaml.safe_dump(yaml_dict, temp_file)

        yield temp_file.name

    finally:

        os.unlink(temp_file.name)

class EnvJujuClient:

    # The environments.yaml options that are replaced by bootstrap options.

    #

    # As described in bug #1538735, default-series and --bootstrap-series must

    # match.  'default-series' should be here, but is omitted so that

    # default-series is always forced to match --bootstrap-series.

    bootstrap_replaces = frozenset(['agent-version'])

    # What feature flags have existed that CI used.

    known_feature_flags = frozenset([

        'actions', 'jes', 'address-allocation', 'cloudsigma'])

    # What feature flags are used by this version of the juju client.

    used_feature_flags = frozenset(['address-allocation'])

    _show_status = 'show-status'

    @classmethod

    def get_version(cls, juju_path=None):

        if juju_path is None:

            juju_path = 'juju'

        return subprocess.check_output((juju_path, '--version')).strip()

    def enable_feature(self, flag):

        """Enable juju feature by setting the given flag.

        New versions of juju with the feature enabled by default will silently

        allow this call, but will not export the environment variable.

        """

        if flag not in self.known_feature_flags:

            raise ValueError('Unknown feature flag: %r' % (flag,))

        self.feature_flags.add(flag)

    def get_jes_command(self):

        """For Juju 2.0, this is always kill-controller."""

        return KILL_CONTROLLER

    def is_jes_enabled(self):

        """Does the state-server support multiple environments."""

        try:

            self.get_jes_command()

            return True

        except JESNotSupported:

            return False

    def enable_jes(self):

        """Enable JES if JES is optional.

        Specifically implemented by the clients that optionally support JES.

        This version raises either JESByDefault or JESNotSupported.

        :raises: JESByDefault when JES is always enabled; Juju has the

            'destroy-controller' command.

        :raises: JESNotSupported when JES is not supported; Juju does not have

            the 'system kill' command when the JES feature flag is set.

        """

        if self.is_jes_enabled():

            raise JESByDefault()

        else:

            raise JESNotSupported()

    @classmethod

    def get_full_path(cls):

        if sys.platform == 'win32':

            return WIN_JUJU_CMD

        return subprocess.check_output(('which', 'juju')).rstrip('\n')

    @classmethod

    def by_version(cls, env, juju_path=None, debug=False):

        version = cls.get_version(juju_path)

        if juju_path is None:

            full_path = cls.get_full_path()

        else:

            full_path = os.path.abspath(juju_path)

        if version.startswith('1.16'):

            raise Exception('Unsupported juju: %s' % version)

        elif re.match('^1\.22[.-]', version):

            client_class = EnvJujuClient22

        elif re.match('^1\.24[.-]', version):

            client_class = EnvJujuClient24

        elif re.match('^1\.25[.-]', version):

            client_class = EnvJujuClient25

        elif re.match('^1\.26[.-]', version):

            client_class = EnvJujuClient26

        elif re.match('^1\.', version):

            client_class = EnvJujuClient1X

        elif re.match('^2\.0-alpha1', version):

            client_class = EnvJujuClient2A1

        elif re.match('^2\.0-alpha2', version):

            client_class = EnvJujuClient2A2

        elif re.match('^2\.0-(alpha3|beta[12])', version):

            client_class = EnvJujuClient2B2

        else:

            client_class = EnvJujuClient

        return client_class(env, version, full_path, debug=debug)

    def clone(self, env=None, version=None, full_path=None, debug=None,

              cls=None):

        """Create a clone of this EnvJujuClient.

        By default, the class, environment, version, full_path, and debug

        settings will match the original, but each can be overridden.

        """

        if env is None:

            env = self.env

        if version is None:

            version = self.version

        if full_path is None:

            full_path = self.full_path

        if debug is None:

            debug = self.debug

        if cls is None:

            cls = self.__class__

        other = cls(env, version, full_path, debug=debug)

        other.feature_flags.update(

            self.feature_flags.intersection(other.used_feature_flags))

        return other

    def get_cache_path(self):

        return get_cache_path(self.env.juju_home, models=True)

    def _full_args(self, command, sudo, args,

                   timeout=None, include_e=True, admin=False):

        # sudo is not needed for devel releases.

        if admin:

            e_arg = ('-m', self.get_admin_model_name())

        elif self.env is None or not include_e:

            e_arg = ()

        else:

            e_arg = ('-m', self.env.environment)

        if timeout is None:

            prefix = ()

        else:

            prefix = get_timeout_prefix(timeout, self._timeout_path)

        logging = '--debug' if self.debug else '--show-log'

        # If args is a string, make it a tuple. This makes writing commands

        # with one argument a bit nicer.

        if isinstance(args, basestring):

            args = (args,)

        # we split the command here so that the caller can control where the -e

        # <env> flag goes.  Everything in the command string is put before the

        # -e flag.

        command = command.split()

        return prefix + ('juju', logging,) + tuple(command) + e_arg + args

    @staticmethod

    def _get_env(env):

        if not isinstance(env, JujuData) and isinstance(env,

                                                        SimpleEnvironment):

            # FIXME: JujuData should be used from the start.

            env = JujuData.from_env(env)

        return env

    def __init__(self, env, version, full_path, juju_home=None, debug=False):

        self.env = self._get_env(env)

        self.version = version

        self.full_path = full_path

        self.debug = debug

        self.feature_flags = set()

        if env is not None:

            if juju_home is None:

                if env.juju_home is None:

                    env.juju_home = get_juju_home()

            else:

                env.juju_home = juju_home

        self.juju_timings = {}

        self._timeout_path = get_timeout_path()

    def _shell_environ(self):

        """Generate a suitable shell environment.

        Juju's directory must be in the PATH to support plugins.

        """

        env = dict(os.environ)

        if self.full_path is not None:

            env['PATH'] = '{}{}{}'.format(os.path.dirname(self.full_path),

                                          os.pathsep, env['PATH'])

        flags = self.feature_flags.intersection(self.used_feature_flags)

        if flags:

            env[JUJU_DEV_FEATURE_FLAGS] = ','.join(sorted(flags))

        env['JUJU_DATA'] = self.env.juju_home

        return env

    def add_ssh_machines(self, machines):

        for count, machine in enumerate(machines):

            try:

                self.juju('add-machine', ('ssh:' + machine,))

            except subprocess.CalledProcessError:

                if count != 0:

                    raise

                logging.warning('add-machine failed.  Will retry.')

                pause(30)

                self.juju('add-machine', ('ssh:' + machine,))

    @staticmethod

    def get_cloud_region(cloud, region):

        if region is None:

            return cloud

        return '{}/{}'.format(cloud, region)

    def get_bootstrap_args(self, upload_tools, config_filename,

                           bootstrap_series=None):

        """Return the bootstrap arguments for the substrate."""

        if self.env.maas:

            constraints = 'mem=2G arch=amd64'

        elif self.env.joyent:

            # Only accept kvm packages by requiring >1 cpu core, see lp:1446264

            constraints = 'mem=2G cpu-cores=1'

        else:

            constraints = 'mem=2G'

        cloud_region = self.get_cloud_region(self.env.get_cloud(),

                                             self.env.get_region())

        args = ['--constraints', constraints, self.env.environment,

                cloud_region, '--config', config_filename,

                '--default-model', self.env.environment]

        if upload_tools:

            args.insert(0, '--upload-tools')

        else:

            args.extend(['--agent-version', self.get_matching_agent_version()])

        if bootstrap_series is not None:

            args.extend(['--bootstrap-series', bootstrap_series])

        return tuple(args)

    @contextmanager

    def _bootstrap_config(self):

        config_dict = make_safe_config(self)

        # Strip unneeded variables.

        config_dict = dict((k, v) for k, v in config_dict.items() if k not in {

            'access-key',

            'admin-secret',

            'application-id',

            'application-password',

            'auth-url',

            'bootstrap-host',

            'client-email',

            'client-id',

            'control-bucket',

            'location',

            'maas-oauth',

            'maas-server',

            'manta-key-id',

            'manta-user',

            'name',

            'password',

            'private-key',

            'region',

            'sdc-key-id',

            'sdc-url',

            'sdc-user',

            'secret-key',

            'storage-account-name',

            'subscription-id',

            'tenant-id',

            'tenant-name',

            'type',

            'username',

        })

        with temp_yaml_file(config_dict) as config_filename:

            yield config_filename

    def _check_bootstrap(self):

        if self.env.environment != self.env.controller.name:

            raise AssertionError(

                'Controller and environment names should not vary (yet)')

    def bootstrap(self, upload_tools=False, bootstrap_series=None):

        """Bootstrap a controller."""

        self._check_bootstrap()

        with self._bootstrap_config() as config_filename:

            args = self.get_bootstrap_args(

                upload_tools, config_filename, bootstrap_series)

            self.juju('bootstrap', args, include_e=False)

    @contextmanager

    def bootstrap_async(self, upload_tools=False, bootstrap_series=None):

        self._check_bootstrap()

        with self._bootstrap_config() as config_filename:

            args = self.get_bootstrap_args(

                upload_tools, config_filename, bootstrap_series)

            with self.juju_async('bootstrap', args, include_e=False):

                yield

                log.info('Waiting for bootstrap of {}.'.format(

                    self.env.environment))

    def create_environment(self, controller_client, config_file):

        controller_client.controller_juju('create-model', (

            self.env.environment, '--config', config_file))

    def destroy_model(self):

        exit_status = self.juju(

            'destroy-model', (self.env.environment, '-y',),

            include_e=False, timeout=timedelta(minutes=10).total_seconds())

        return exit_status

    def kill_controller(self):

        """Kill a controller and its environments."""

        seen_cmd = self.get_jes_command()

        self.juju(

            _jes_cmds[seen_cmd]['kill'], (self.env.controller.name, '-y'),

            include_e=False, check=False, timeout=600)

    def get_juju_output(self, command, *args, **kwargs):

        """Call a juju command and return the output.

        Sub process will be called as 'juju <command> <args> <kwargs>'. Note

        that <command> may be a space delimited list of arguments. The -e

        <environment> flag will be placed after <command> and before args.

        """

        args = self._full_args(command, False, args,

                               timeout=kwargs.get('timeout'),

                               include_e=kwargs.get('include_e', True),

                               admin=kwargs.get('admin', False))

        env = self._shell_environ()

        log.debug(args)

        # Mutate os.environ instead of supplying env parameter so

        # Windows can search env['PATH']

        with scoped_environ(env):

            proc = subprocess.Popen(

                args, stdout=subprocess.PIPE, stdin=subprocess.PIPE,

                stderr=subprocess.PIPE)

            sub_output, sub_error = proc.communicate()

            log.debug(sub_output)

            if proc.returncode != 0:

                log.debug(sub_error)

                e = subprocess.CalledProcessError(

                    proc.returncode, args, sub_output)

                e.stderr = sub_error

                if (

                    'Unable to connect to environment' in sub_error or

                        'MissingOrIncorrectVersionHeader' in sub_error or

                        '307: Temporary Redirect' in sub_error):

                    raise CannotConnectEnv(e)

                raise e

        return sub_output

    def show_status(self):

        """Print the status to output."""

        self.juju(self._show_status, ('--format', 'yaml'))

    def get_status(self, timeout=60, raw=False, admin=False, *args):

        """Get the current status as a dict."""

        # GZ 2015-12-16: Pass remaining timeout into get_juju_output call.

        for ignored in until_timeout(timeout):

            try:

                if raw:

                    return self.get_juju_output(self._show_status, *args)

                return Status.from_text(

                    self.get_juju_output(

                        self._show_status, '--format', 'yaml', admin=admin))

            except subprocess.CalledProcessError:

                pass

        raise Exception(

            'Timed out waiting for juju status to succeed')

    @staticmethod

    def _dict_as_option_strings(options):

        return tuple('{}={}'.format(*item) for item in options.items())

    def set_config(self, service, options):

        option_strings = self._dict_as_option_strings(options)

        self.juju('set-config', (service,) + option_strings)

    def get_config(self, service):

        return yaml_loads(self.get_juju_output('get-config', service))

    def get_service_config(self, service, timeout=60):

        for ignored in until_timeout(timeout):

            try:

                return self.get_config(service)

            except subprocess.CalledProcessError:

                pass

        raise Exception(

            'Timed out waiting for juju get %s' % (service))

    def set_model_constraints(self, constraints):

        constraint_strings = self._dict_as_option_strings(constraints)

        return self.juju('set-model-constraints', constraint_strings)

    def get_model_config(self):

        """Return the value of the environment's configured option."""

        return yaml.safe_load(self.get_juju_output('get-model-config'))

    def get_env_option(self, option):

        """Return the value of the environment's configured option."""

        return self.get_juju_output('get-model-config', option)

    def set_env_option(self, option, value):

        """Set the value of the option in the environment."""

        option_value = "%s=%s" % (option, value)

        return self.juju('set-model-config', (option_value,))

    def set_testing_tools_metadata_url(self):

        url = self.get_env_option('tools-metadata-url')

        if 'testing' not in url:

            testing_url = url.replace('/tools', '/testing/tools')

            self.set_env_option('tools-metadata-url', testing_url)

    def juju(self, command, args, sudo=False, check=True, include_e=True,

             timeout=None, extra_env=None):

        """Run a command under juju for the current environment."""

        args = self._full_args(command, sudo, args, include_e=include_e,

                               timeout=timeout)

        log.info(' '.join(args))

        env = self._shell_environ()

        if extra_env is not None:

            env.update(extra_env)

        if check:

            call_func = subprocess.check_call

        else:

            call_func = subprocess.call

        start_time = time.time()

        # Mutate os.environ instead of supplying env parameter so Windows can

        # search env['PATH']

        with scoped_environ(env):

            rval = call_func(args)

        self.juju_timings.setdefault(args, []).append(

            (time.time() - start_time))