<commit_before>import os
import sys
import logging
import json
import random
import string
from datetime import datetime
from importlib import import_module
from path import path
import yaml
from kv import KV
import pkg_resources
from .daemons import Supervisor
from .routing import Haproxy
from . import deployer


log = logging.getLogger(__name__)


CFG_LINKS_FOLDER = 'active'
YAML_EXT = '.yaml'
RUN_RC_NAME = '.runrc'


def _get_named_object(name):
    module_name, attr_name = name.split(':')
    module = import_module(module_name)
    return getattr(module, attr_name)


def random_id(size=6, vocabulary=string.ascii_lowercase + string.digits):
    return ''.join(random.choice(vocabulary) for c in range(size))


class Bucket(object):

    def __init__(self, id_, airship, config):
        self.id_ = id_
        self.airship = airship
        self.config = config
        self.folder = self.airship._bucket_folder(id_)
        var = self.airship.home_path / 'var'
        self.run_folder = var / 'run' / id_
        self.log_path = var / 'log' / (self.id_ + '.log')

    @property
    def meta(self):
        return self.config['meta']

    @property
    def port(self):
        return self.config['port']

    def configure(self):
        self.run_folder.makedirs_p()

    def start(self):
        log.info("Activating bucket %r", self.id_)
        self.configure()
        self.airship.daemons.configure_bucket_running(self)
        self.airship.haproxy.configure_bucket(self)

    def stop(self):
        self.airship.haproxy.remove_bucket(self)
        self.airship.daemons.configure_bucket_stopped(self)

    def trigger(self):
        self.airship.daemons.trigger_bucket(self)

    def destroy(self):
        self.airship.daemons.remove_bucket(self.id_)
        if self.run_folder.isdir():
            self.run_folder.rmtree()
        if self.folder.isdir():
            self.folder.rmtree()
        self.airship.buckets_db.pop(self.id_, None)
        self.airship._free_port(self)

    def run(self, command):
        os.chdir(self.folder)
        environ = dict(os.environ)
        environ.update(self.airship.config.get('env') or {})
        environ['PORT'] = str(self.port)
        venv = self.folder / '_virtualenv'
        if venv.isdir():
            environ['PATH'] = ((venv / 'bin') + ':' + environ['PATH'])
        shell_args = ['/bin/bash']
        if command:
            environ['BASH_ENV'] = RUN_RC_NAME
            shell_args += ['-c', command]
        else:
            shell_args += ['--rcfile', RUN_RC_NAME]
        os.execve(shell_args[0], shell_args, environ)


class Airship(object):
    """ The airship object implements most operations performed by airship. It
    acts as container for deployments.
    """

    def __init__(self, config):
        self.home_path = config['home']
        self.config = config
        etc = self.home_path / 'etc'
        etc.mkdir_p()
        self.buckets_db = KV(etc / 'buckets.db', table='bucket')
        self.daemons = Supervisor(etc)
        self.haproxy = Haproxy(self.home_path, config.get('port_map') or {})
        from routing import configuration_update
        configuration_update.connect(self._haproxy_update, self.haproxy)

    @property
    def cfg_links_folder(self):
        folder = self.home_path / CFG_LINKS_FOLDER
        if not folder.isdir():
            folder.makedirs()
        return folder

    def _haproxy_configure_supervisor(self):
        haproxy_program = self.home_path / 'etc' / 'supervisor.d' / 'haproxy'
        haproxy_program.write_text(self.haproxy.supervisord_config(self))

    def initialize(self):
        self.generate_supervisord_configuration()
        self._haproxy_configure_supervisor()

    def _haproxy_update(self, sender, **extra):
        self._haproxy_configure_supervisor()
        self.daemons.ctl(['update'])
        self.daemons.ctl(['restart', 'haproxy'])

    def generate_supervisord_configuration(self):
        self.daemons.configure(self.home_path)

    def _get_bucket_by_id(self, bucket_id):
        config = self.buckets_db[bucket_id]
        return Bucket(bucket_id, self, config)

    def get_bucket(self, name):
        try:
            return self._get_bucket_by_id(name)

        except KeyError:
            for bucket_id in self.buckets_db:
                bucket = self._get_bucket_by_id(bucket_id)
                if name == bucket.meta.get('APPLICATION_NAME'):
                    return bucket

            else:
                raise KeyError

    def _bucket_folder(self, id_):
        return self.home_path / id_

    def _generate_bucket_id(self, id_prefix):
        for c in range(10):
            id_ = id_prefix + random_id()
            try:
                self._bucket_folder(id_).mkdir()
            except OSError:
                continue
            else:
                return id_
        else:
            raise RuntimeError("Failed to generate unique bucket ID")

    def _open_ports_db(self):
        return KV(self.home_path / 'etc' / 'buckets.db', table='port')

    def _allocate_port(self, bucket_id):
        from itertools import chain
        port_range = self.config.get('port_range', [5000, 5099])
        start_port = port_range[0]
        end_port = port_range[1] + 1

        ports_db = self._open_ports_db()
        with ports_db.lock():
            next_port = ports_db.get('next', start_port)
            if start_port <= next_port <= end_port:
                queue = chain(xrange(next_port, end_port),
                              xrange(start_port, next_port - 1))
            else:
                queue = xrange(start_port, end_port)
            for port in queue:
                assert start_port <= port <= end_port
                if port not in ports_db:
                    ports_db[port] = bucket_id
                    ports_db['next'] = port + 1
                    return port
            else:
                raise RuntimeError("No ports free to allocate")

    def _free_port(self, bucket):
        ports_db = self._open_ports_db()
        port = ports_db.pop(bucket.port, None)
        if port is not None:
            assert port == bucket.id_

    def new_bucket(self, config={}):
        meta = {'CREATION_TIME': datetime.utcnow().isoformat()}
        app_name = config.get('application_name')
        if app_name:
            meta['APPLICATION_NAME'] = app_name
            id_prefix = app_name + '-'
        else:
            id_prefix = ''
        bucket_id = self._generate_bucket_id(id_prefix)
        self.buckets_db[bucket_id] = {
            'require-services': config.get('services', {}),
            'urlmap': config.get('urlmap', []),
            'meta': meta,
            'port': self._allocate_port(bucket_id),
        }
        bucket = self._get_bucket_by_id(bucket_id)
        return bucket

    def list_buckets(self):
        buckets = []
        for bucket_id in self.buckets_db:
            bucket = self._get_bucket_by_id(bucket_id)
            buckets.append({
                'id': bucket.id_,
                'meta': bucket.meta,
                'port': bucket.port,
            })
        return {'buckets': buckets}


def load_plugins():
    for entry_point in pkg_resources.iter_entry_points('sarge_plugins'):
        callback = entry_point.load()
        callback()


AIRSHIP_SCRIPT = """#!/bin/bash
exec '{prefix}/bin/airship' '{home}' "$@"
"""

SUPERVISORD_SCRIPT = """#!/bin/bash
exec '{prefix}/bin/supervisord' -c '{home}/etc/supervisor.conf'
"""

SUPERVISORCTL_SCRIPT = """#!/bin/bash<commit_msg>rename pkg_resources entry point for plugins<commit_after>import os
import sys
import logging
import json
import random
import string
from datetime import datetime
from importlib import import_module
from path import path
import yaml
from kv import KV
import pkg_resources
from .daemons import Supervisor
from .routing import Haproxy
from . import deployer


log = logging.getLogger(__name__)


CFG_LINKS_FOLDER = 'active'
YAML_EXT = '.yaml'
RUN_RC_NAME = '.runrc'


def _get_named_object(name):
    module_name, attr_name = name.split(':')
    module = import_module(module_name)
    return getattr(module, attr_name)


def random_id(size=6, vocabulary=string.ascii_lowercase + string.digits):
    return ''.join(random.choice(vocabulary) for c in range(size))


class Bucket(object):

    def __init__(self, id_, airship, config):
        self.id_ = id_
        self.airship = airship
        self.config = config
        self.folder = self.airship._bucket_folder(id_)
        var = self.airship.home_path / 'var'
        self.run_folder = var / 'run' / id_
        self.log_path = var / 'log' / (self.id_ + '.log')

    @property
    def meta(self):
        return self.config['meta']

    @property
    def port(self):
        return self.config['port']

    def configure(self):
        self.run_folder.makedirs_p()

    def start(self):
        log.info("Activating bucket %r", self.id_)
        self.configure()
        self.airship.daemons.configure_bucket_running(self)
        self.airship.haproxy.configure_bucket(self)

    def stop(self):
        self.airship.haproxy.remove_bucket(self)
        self.airship.daemons.configure_bucket_stopped(self)

    def trigger(self):
        self.airship.daemons.trigger_bucket(self)

    def destroy(self):
        self.airship.daemons.remove_bucket(self.id_)
        if self.run_folder.isdir():
            self.run_folder.rmtree()
        if self.folder.isdir():
            self.folder.rmtree()
        self.airship.buckets_db.pop(self.id_, None)
        self.airship._free_port(self)

    def run(self, command):
        os.chdir(self.folder)
        environ = dict(os.environ)
        environ.update(self.airship.config.get('env') or {})
        environ['PORT'] = str(self.port)
        venv = self.folder / '_virtualenv'
        if venv.isdir():
            environ['PATH'] = ((venv / 'bin') + ':' + environ['PATH'])
        shell_args = ['/bin/bash']
        if command:
            environ['BASH_ENV'] = RUN_RC_NAME
            shell_args += ['-c', command]
        else:
            shell_args += ['--rcfile', RUN_RC_NAME]
        os.execve(shell_args[0], shell_args, environ)


class Airship(object):
    """ The airship object implements most operations performed by airship. It
    acts as container for deployments.
    """

    def __init__(self, config):
        self.home_path = config['home']
        self.config = config
        etc = self.home_path / 'etc'
        etc.mkdir_p()
        self.buckets_db = KV(etc / 'buckets.db', table='bucket')
        self.daemons = Supervisor(etc)
        self.haproxy = Haproxy(self.home_path, config.get('port_map') or {})
        from routing import configuration_update
        configuration_update.connect(self._haproxy_update, self.haproxy)

    @property
    def cfg_links_folder(self):
        folder = self.home_path / CFG_LINKS_FOLDER
        if not folder.isdir():
            folder.makedirs()
        return folder

    def _haproxy_configure_supervisor(self):
        haproxy_program = self.home_path / 'etc' / 'supervisor.d' / 'haproxy'
        haproxy_program.write_text(self.haproxy.supervisord_config(self))

    def initialize(self):
        self.generate_supervisord_configuration()
        self._haproxy_configure_supervisor()

    def _haproxy_update(self, sender, **extra):
        self._haproxy_configure_supervisor()
        self.daemons.ctl(['update'])
        self.daemons.ctl(['restart', 'haproxy'])

    def generate_supervisord_configuration(self):
        self.daemons.configure(self.home_path)

    def _get_bucket_by_id(self, bucket_id):
        config = self.buckets_db[bucket_id]
        return Bucket(bucket_id, self, config)

    def get_bucket(self, name):
        try:
            return self._get_bucket_by_id(name)

        except KeyError:
            for bucket_id in self.buckets_db:
                bucket = self._get_bucket_by_id(bucket_id)
                if name == bucket.meta.get('APPLICATION_NAME'):
                    return bucket

            else:
                raise KeyError

    def _bucket_folder(self, id_):
        return self.home_path / id_

    def _generate_bucket_id(self, id_prefix):
        for c in range(10):
            id_ = id_prefix + random_id()
            try:
                self._bucket_folder(id_).mkdir()
            except OSError:
                continue
            else:
                return id_
        else:
            raise RuntimeError("Failed to generate unique bucket ID")

    def _open_ports_db(self):
        return KV(self.home_path / 'etc' / 'buckets.db', table='port')

    def _allocate_port(self, bucket_id):
        from itertools import chain
        port_range = self.config.get('port_range', [5000, 5099])
        start_port = port_range[0]
        end_port = port_range[1] + 1

        ports_db = self._open_ports_db()
        with ports_db.lock():
            next_port = ports_db.get('next', start_port)
            if start_port <= next_port <= end_port:
                queue = chain(xrange(next_port, end_port),
                              xrange(start_port, next_port - 1))
            else:
                queue = xrange(start_port, end_port)
            for port in queue:
                assert start_port <= port <= end_port
                if port not in ports_db:
                    ports_db[port] = bucket_id
                    ports_db['next'] = port + 1
                    return port
            else:
                raise RuntimeError("No ports free to allocate")

    def _free_port(self, bucket):
        ports_db = self._open_ports_db()
        port = ports_db.pop(bucket.port, None)
        if port is not None:
            assert port == bucket.id_

    def new_bucket(self, config={}):
        meta = {'CREATION_TIME': datetime.utcnow().isoformat()}
        app_name = config.get('application_name')
        if app_name:
            meta['APPLICATION_NAME'] = app_name
            id_prefix = app_name + '-'
        else:
            id_prefix = ''
        bucket_id = self._generate_bucket_id(id_prefix)
        self.buckets_db[bucket_id] = {
            'require-services': config.get('services', {}),
            'urlmap': config.get('urlmap', []),
            'meta': meta,
            'port': self._allocate_port(bucket_id),
        }
        bucket = self._get_bucket_by_id(bucket_id)
        return bucket

    def list_buckets(self):
        buckets = []
        for bucket_id in self.buckets_db:
            bucket = self._get_bucket_by_id(bucket_id)
            buckets.append({
                'id': bucket.id_,
                'meta': bucket.meta,
                'port': bucket.port,
            })
        return {'buckets': buckets}


def load_plugins():
    for entry_point in pkg_resources.iter_entry_points('airship_plugins'):
        callback = entry_point.load()
        callback()


AIRSHIP_SCRIPT = """#!/bin/bash
exec '{prefix}/bin/airship' '{home}' "$@"
"""

SUPERVISORD_SCRIPT = """#!/bin/bash
exec '{prefix}/bin/supervisord' -c '{home}/etc/supervisor.conf'
"""

SUPERVISORCTL_SCRIPT = """#!/bin/bash
{'fraction_non_alphanumeric': 0.111162575266327, 'fraction_numerical': 0.0022497187851518562, 'mean_word_length': 2.258732212160414, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5346247', 'n_tokens_mistral': 4823, 'n_tokens_neox': 4646, 'n_words': 954}