juju-sync-watch

#!/usr/bin/env python

import re
import os
import sys
import yaml
import subprocess
import argparse
from datetime import datetime

import jujuclient
import jujuclient.juju1.environment
import jujuclient.juju2.environment


def call(*args):
     try:
         return subprocess.check_output(('/usr/bin/env',) + args,
                                        stderr=subprocess.STDOUT)
     except subprocess.CalledProcessError as e:
         sys.stderr.write(e.output)
         sys.exit(e.returncode)


JUJU_REPOSITORY = os.environ.get('JUJU_REPOSITORY', '')
JUJU_ENV = call('juju', 'switch').rstrip('\n')
JUJU_VERSION = call('juju', 'version')
JUJU_2 = JUJU_VERSION.startswith('2.')


class ShowDesc(argparse.Action):
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        super(ShowDesc, self).__init__(option_strings, dest, nargs=0, **kwargs)

    def __call__(*args, **kwargs):
        print 'Automatic syncing of local changes to deployed service'
        sys.exit(0)


def parse_args():
    parser = argparse.ArgumentParser(
        description='Automatic syncing of local changes to deployed service',
        epilog="""
            This will watch for changes in the local source of a deployed charm
            and automatically push the changes up to the remote unit and retry
            the last failed hook.  If CHARM_PATH is not given, it will look first in
            the current directory and then in $JUJU_REPOSITORY/$SERIES/$CHARM for the
            charm source.

            Note: This requires pyinotify to be installed (pip install pyinotify).
        """)
    parser.add_argument('unit')
    parser.add_argument('charm_path', nargs='?')
    parser.add_argument('-d', '--description', action=ShowDesc, help='show description')
    parser.add_argument('-e', '--environment', dest='env', default=JUJU_ENV,
                        help='Juju environment or model name. Defaults to the value '
                             'obtained from `juju switch`')
    parser.add_argument('-R', '--no-retry', '--no-restart', dest='retry', action='store_false',
                        help='do not automatically retry failed hooks when a change is made')
    parser.add_argument('-q', '--quiet', action='store_true', help='suppress output')
    opts = parser.parse_args()
    return opts


def fail(msg):
    sys.stderr.write('%s\n' % msg)
    sys.exit(1)


def get_env(env_name=None):
    if hasattr(get_env, '_env'):
        return get_env._env
    if JUJU_2:
        get_env._env = jujuclient.juju2.environment.Environment.connect(opts.env)
    else:
        get_env._env = jujuclient.juju1.environment.Environment.connect(opts.env)
    return get_env._env


def parse_charm_id(charm_id):
    charm_match = re.match(r'local:([^/]*)/(.*)-\d+', charm_id)
    if charm_match:
        return charm_match.groups()
    return None, None


def run(*args):
    try:
        return subprocess.check_output(args, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        if 'already resolved' in e.output:
            return e.output
        else:
            fail(e.output)


def watch(unit_name, charm_path, retry, quiet):
    import pyinotify
    wm = pyinotify.WatchManager()
    mask = pyinotify.IN_DELETE | pyinotify.IN_CREATE | pyinotify.IN_MODIFY

    class EventHandler(pyinotify.ProcessEvent):
        last_event = None
        changes = False

        def process_default(self, event):
            self.last_event = datetime.now()

        def dispatch_update(self):
            if self.last_event:
                self.changes = True
                if not quiet:
                    print 'Updating %s...' % unit_name,
                    sys.stdout.flush()
                update_unit(unit_name, charm_path)
                if not quiet:
                    print 'done'
                if retry and unit_in_error(unit_name):
                    if not quiet:
                        print 'Retrying failed hook'
                    restart_errored_unit(unit_name)
                self.last_event = None

    handler = EventHandler()
    notifier = pyinotify.Notifier(wm, handler, timeout=2000)
    wm.add_watch(charm_path, mask, rec=True)
    if not quiet:
        print 'Watching %s' % charm_path
    try:
        while True:
            notifier.process_events()
            while notifier.check_events():  # process events until a timeout
                notifier.read_events()
                notifier.process_events()
            handler.dispatch_update()
    except KeyboardInterrupt:
        return handler.changes


def update_unit(unit_name, charm_path):
    remote_path = '/var/lib/juju/agents/unit-%s/charm/' % unit_name.replace('/', '-')
    run('juju', 'scp', '--', '-r', charm_path, '%s:.juju-sync-watch' % unit_name)
    run('juju', 'ssh', unit_name, ';'.join(map('sudo {}'.format, [
        'chown -R root:root .juju-sync-watch',
        'cp -R .juju-sync-watch/* %s' % remote_path,
        'rm -rf .juju-sync-watch'])))


def unit_in_error(unit_name):
    env = get_env()
    service_name = unit_name.split('/')[0]
    status = env.status()
    if JUJU_2:
        service = status['applications'][service_name]
        if service.get('subordinate-to'):
            service = status['applications'][service['subordinate-to'][0]]
            unit = filter(None, [unit['subordinates'].get(unit_name) for unit in service['units'].values()])[0]
        else:
            unit = service['units'][unit_name]
        unit_status = unit.get('workload-status', {}).get('status')
    else:
        service = status['Services'][service_name]
        if service.get('SubordinateTo'):
            service = status['Services'][service['SubordinateTo'][0]]
            unit = filter(None, [unit['Subordinates'].get(unit_name) for unit in service['Units'].values()])[0]
        else:
            unit = service['Units'][unit_name]
        unit_status = unit.get('AgentState')

    return unit_status == 'error'


def restart_errored_unit(unit_name):
    if JUJU_2:
        run('juju', 'resolved', unit_name)
    else:
        run('juju', 'resolved', '--retry', unit_name)


def is_charm_dir(charm_name, charm_path):
    metadata_yaml = os.path.join(charm_path, 'metadata.yaml')
    if not os.path.exists(metadata_yaml):
        return False
    with open(metadata_yaml) as fp:
        metadata = yaml.safe_load(fp)
        return metadata['name'] == charm_name


def get_charm_info(env, unit_name):
    if '/' not in unit_name:
        fail('Invalid unit name: %s' % unit_name)
    service_name = unit_name.split('/')[0]
    status = env.status()
    services = status['Services']
    if service_name not in services:
        fail('Unit not found: %s' % unit_name)
    service = services[service_name]
    series, charm_name = parse_charm_id(service['Charm'])
    if not series or not charm_name:
        fail('Unable to determine series or charm_name')
    return series, charm_name


def find_charm(series, charm_name):
    candidates = ['.']
    if JUJU_REPOSITORY:
        candidates.append('%s/%s/%s' % (JUJU_REPOSITORY, series, charm_name))
    else:
        candidates.append('%s/%s' % (series, charm_name))
    for charm_path in candidates:
        if is_charm_dir(charm_name, charm_path):
            return charm_path
    fail('Unable to find source path for %s' % charm_name)


def prompt_for_upgrade(unit_name):
    service_name = unit_name.split('/')[0]
    print
    response = raw_input('Upgrade charm %s? [y/N] ' % service_name)
    if response.lower().startswith('y'):
        print run('juju', 'upgrade-charm', service_name)


if __name__ == '__main__':
    opts = parse_args()
    env = get_env(opts.env)
    if not hasattr(opts, 'charm_path'):
        series, charm_name = get_charm_info(env, opts.unit)
        opts.charm_path = find_charm(series, charm_name)
    changes = watch(opts.unit, opts.charm_path, opts.retry, opts.quiet)

    # this doesn't work with Juju <= 1.23 due to way plugins are subprocessed :(
    #if changes:
    #    prompt_for_upgrade(opts.unit)