upgrade

#!/usr/bin/env python3
import urllib
import argparse
import os
import subprocess
import sys
from collections import namedtuple
from fnmatch import fnmatch

import requests

try:
    import semver
except ImportError:
    print("Missing required library: semver.")
    exit(1)

REDASH_HOME = os.environ.get('REDASH_HOME', '/opt/redash')
CURRENT_VERSION_PATH = '{}/current'.format(REDASH_HOME)


def run(cmd, cwd=None):
    if not cwd:
        cwd = REDASH_HOME

    return subprocess.check_output(cmd, cwd=cwd, shell=True, stderr=subprocess.STDOUT)


def confirm(question):
    reply = str(input(question + ' (y/n): ')).lower().strip()

    if reply[0] == 'y':
        return True
    if reply[0] == 'n':
        return False
    else:
        return confirm("Please use 'y' or 'n'")


def version_path(version_name):
    return "{}/{}".format(REDASH_HOME, version_name)

END_CODE = '\033[0m'


def colored_string(text, color):
    if sys.stdout.isatty():
        return "{}{}{}".format(color, text, END_CODE)
    else:
        return text


def h1(text):
    print(colored_string(text, '\033[4m\033[1m'))


def green(text):
    print(colored_string(text, '\033[92m'))


def red(text):
    print(colored_string(text, '\033[91m'))


class Release(namedtuple('Release', ('version', 'download_url', 'filename', 'description'))):
    def v1_or_newer(self):
        return semver.compare(self.version, '1.0.0-alpha') >= 0

    def is_newer(self, version):
        return semver.compare(self.version, version) > 0

    @property
    def version_name(self):
        return self.filename.replace('.tar.gz', '')


def get_latest_release_from_ci():
    response = requests.get('https://circleci.com/api/v1.1/project/github/getredash/redash/latest/artifacts?branch=master')

    if response.status_code != 200:
        exit("Failed getting releases (status code: %s)." % response.status_code)

    tarball_asset = filter(lambda asset: asset['url'].endswith('.tar.gz'), response.json())[0]
    filename = urllib.unquote(tarball_asset['pretty_path'].split('/')[-1])
    version = filename.replace('redash.', '').replace('.tar.gz', '')

    release = Release(version, tarball_asset['url'], filename, '')

    return release


def get_release(channel):
    if channel == 'ci':
        return get_latest_release_from_ci()

    response = requests.get('https://version.redash.io/api/releases?channel={}'.format(channel))
    release = response.json()[0]

    filename = release['download_url'].split('/')[-1]
    release = Release(release['version'], release['download_url'], filename, release['description'])

    return release


def link_to_current(version_name):
    green("Linking to current version...")
    run('ln -nfs {} {}'.format(version_path(version_name), CURRENT_VERSION_PATH))


def restart_services():
    # We're doing this instead of simple 'supervisorctl restart all' because
    # otherwise it won't notice that /opt/redash/current pointing at a different
    # directory.
    green("Restarting...")
    try:
        run('sudo /etc/init.d/redash_supervisord restart')
    except subprocess.CalledProcessError as e:
        run('sudo service supervisor restart')


def update_requirements(version_name):
    green("Installing new Python packages (if needed)...")
    new_requirements_file = '{}/requirements.txt'.format(version_path(version_name))

    install_requirements = False

    try:
        run('diff {}/requirements.txt {}'.format(CURRENT_VERSION_PATH, new_requirements_file)) != 0
    except subprocess.CalledProcessError as e:
        if e.returncode != 0:
            install_requirements = True

    if install_requirements:
        run('sudo pip install -r {}'.format(new_requirements_file))


def apply_migrations(release):
    green("Running migrations (if needed)...")
    if not release.v1_or_newer():
        return apply_migrations_pre_v1(release.version_name)

    run("sudo -u redash bin/run ./manage.py db upgrade", cwd=version_path(release.version_name))


def find_migrations(version_name):
    current_migrations = set([f for f in os.listdir("{}/migrations".format(CURRENT_VERSION_PATH)) if fnmatch(f, '*_*.py')])
    new_migrations = sorted([f for f in os.listdir("{}/migrations".format(version_path(version_name))) if fnmatch(f, '*_*.py')])

    return [m for m in new_migrations if m not in current_migrations]


def apply_migrations_pre_v1(version_name):
    new_migrations = find_migrations(version_name)

    if new_migrations:
        green("New migrations to run: ")
        print(', '.join(new_migrations))
    else:
        print("No new migrations in this version.")

    if new_migrations and confirm("Apply new migrations? (make sure you have backup)"):
        for migration in new_migrations:
            print("Applying {}...".format(migration))
            run("sudo sudo -u redash PYTHONPATH=. bin/run python migrations/{}".format(migration), cwd=version_path(version_name))


def download_and_unpack(release):
    directory_name = release.version_name

    green("Downloading release tarball...")
    run('sudo wget --header="Accept: application/octet-stream" -O {} {}'.format(release.filename, release.download_url))
    green("Unpacking to: {}...".format(directory_name))
    run('sudo mkdir -p {}'.format(directory_name))
    run('sudo tar -C {} -xvf {}'.format(directory_name, release.filename))

    green("Changing ownership to redash...")
    run('sudo chown redash {}'.format(directory_name))

    green("Linking .env file...")
    run('sudo ln -nfs {}/.env {}/.env'.format(REDASH_HOME, version_path(directory_name)))


def current_version():
    real_current_path = os.path.realpath(CURRENT_VERSION_PATH).replace('.b', '+b')
    return real_current_path.replace(REDASH_HOME + '/', '').replace('redash.', '')


def verify_minimum_version():
    green("Current version: " + current_version())
    if semver.compare(current_version(), '0.12.0') < 0:
        red("You need to have Redash v0.12.0 or newer to upgrade to post v1.0.0 releases.")
        green("To upgrade to v0.12.0, run the upgrade script set to the legacy channel (--channel legacy).")
        exit(1)


def show_description_and_confirm(description):
    if description:
        print(description)

        if not confirm("Continue with upgrade?"):
            red("Cancelling upgrade.")
            exit(1)


def verify_newer_version(release):
    if not release.is_newer(current_version()):
        red("The found release is not newer than your current deployed release ({}).".format(current_version()))
        if not confirm("Continue with upgrade?"):
            red("Cancelling upgrade.")
            exit(1)


def deploy_release(channel):
    h1("Starting Redash upgrade:")

    release = get_release(channel)
    green("Found version: {}".format(release.version))

    if release.v1_or_newer():
        verify_minimum_version()

    verify_newer_version(release)
    show_description_and_confirm(release.description)

    try:
        download_and_unpack(release)
        update_requirements(release.version_name)
        apply_migrations(release)
        link_to_current(release.version_name)
        restart_services()
        green("Done! Enjoy.")
    except subprocess.CalledProcessError as e:
        red("Failed running: {}".format(e.cmd))
        red("Exit status: {}\nOutput:\n{}".format(e.returncode, e.output))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--channel", help="The channel to get release from (default: stable).", default='stable')
    args = parser.parse_args()

    deploy_release(args.channel)