class SliderConfig(BaseConfig):
    '''
    presto-admin needs to update the slider config other than through the
    interactive config process because it needs to keep track of the name
    of the presto-yarn-package we install.

    As a result, SliderConfig acts a little funny; it acts like enough of a
    dict to allow env.conf[NAME] lookups and modifications, and it also exposes
    the ability to store the config after it's been modified.
    '''

    def __init__(self):
        super(SliderConfig, self).__init__(SLIDER_CONFIG_PATH, _SLIDER_CONFIG)

    @overrides
    def is_config_loaded(self):
        return SLIDER_CONFIG_LOADED in env and env[SLIDER_CONFIG_LOADED]

    @overrides
    def set_config_loaded(self):
        env[SLIDER_CONFIG_LOADED] = True

    @overrides
    def set_env_from_conf(self, conf):
        self.config.update(conf)
        env.user = conf[ADMIN_USER]
        env.port = conf[SSH_PORT]
        env.roledefs[SLIDER_MASTER] = [conf[HOST]]
        env.roledefs['all'] = env.roledefs[SLIDER_MASTER]

        env.conf = self

        env.hosts = env.roledefs['all'][:]

    def store_conf(self):
        super(SliderConfig, self).write_conf(self.config) DIR = 'slider_directory' SLIDER_USER = 'slider_user' APPNAME = 'slider_appname' JAVA_HOME = 'JAVA_HOME' HADOOP_CONF = 'HADOOP_CONF' SLIDER_MASTER = 'slider_master' PRESTO_PACKAGE = 'presto_slider_package' SLIDER_CONFIG_DIR = os.path.join(get_config_directory(), 'slider') def requires_config(config_class):
    def wrap(func):
        config_instance = config_class()
        func.pa_config_callback = config_instance.get_config

        @wraps(func)
        def wrapper(*args, **kwargs):
            if not config_instance.is_config_loaded():
                raise ConfigurationError('Required config not loaded at task '
                                         'execution time.')
            return func(*args, **kwargs)
        return wrapper
    return wrap def task_by_rolename(rolename):
    def inner_decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            return by_rolename(env.host, rolename, f, *args, **kwargs)
        return wrapper
    return inner_decorator
import os.path
from fabric.api import env, task, abort
from fabric.context_managers import shell_env
from fabric.operations import put, sudo, local
from prestoadmin.yarn_slider.config import SliderConfig, \
    DIR, SLIDER_USER, APPNAME, JAVA_HOME, HADOOP_CONF, SLIDER_MASTER, \
    PRESTO_PACKAGE, SLIDER_CONFIG_DIR
from prestoadmin.util.base_config import requires_config
from prestoadmin.util.fabricapi import task_by_rolename
# -*- coding: utf-8 -*-
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Module for managing presto/YARN integration.
"""





__all__ = ['install', 'uninstall']


SLIDER_PKG_DEFAULT_FILES = ['appConfig-default.json', 'resources-default.json']


def get_slider_bin(conf):
    return os.path.join(conf[DIR], 'bin', 'slider')


def run_slider(slider_command, conf):
    with shell_env(JAVA_HOME=conf[JAVA_HOME],
                   HADOOP_CONF_DIR=conf[HADOOP_CONF]):
        return sudo(slider_command, user=conf[SLIDER_USER])


@task
@requires_config(SliderConfig)
@task_by_rolename(SLIDER_MASTER)
def install(presto_yarn_package):
    """
    Install the presto-yarn package on the cluster using Apache Slider. The
    presto-yarn package takes the form of a zip file that conforms to Slider's
    packaging requirements. After installing the presto-yarn package the presto
    application is registered with Slider.

    Before Slider can install the presto-yarn package, the slider user's hdfs
    home directory needs to be created. This needs to be done by a user that
    has write access to the hdfs /user directory, typically the user hdfs or a
    member of the superuser group.

    The name of the presto application is arbitrary and set in the slider
    configuration file. The default is PRESTO

    :param presto_yarn_package: The zip file containing the presto-yarn
                                package as structured for Slider.
    """
    conf = env.conf
    package_filename = os.path.basename(presto_yarn_package)
    package_file = os.path.join('/tmp', package_filename)

    result = put(presto_yarn_package, package_file)
    if result.failed:
        abort('Failed to send slider application package to %s on host %s' %
              (package_file, env.host))

    package_install_command = \
        '%s package --install --package %s --name %s' % \
        (get_slider_bin(conf), package_file, conf[APPNAME])

    try:
        run_slider(package_install_command, conf)

        conf[PRESTO_PACKAGE] = package_filename
        conf.store_conf()

        local('unzip %s %s -d %s' %
              (presto_yarn_package, ' '.join(SLIDER_PKG_DEFAULT_FILES),
               SLIDER_CONFIG_DIR))
    finally:
        sudo('rm -f %s' % (package_file))


@task
@requires_config(SliderConfig)
@task_by_rolename(SLIDER_MASTER)
def uninstall():
    """
    Uninstall unregisters the presto application with slider and removes the
    installed package.
    """
    conf = env.conf
    package_delete_command = '%s package --delete --name %s' % \
                             (get_slider_bin(conf), conf[APPNAME])
    run_slider(package_delete_command, conf)

    try:
        del conf[PRESTO_PACKAGE]
        conf.store_conf()
    except KeyError:
        pass


    local('rm %s' % (' '.join([os.path.join(SLIDER_CONFIG_DIR, f)