manage.py

#!/usr/bin/env python
from gevent import monkey; monkey.patch_all()

import copy
import logging
import os
import signal
import subprocess

import click
import gevent
from disco.util.logging import LOG_FORMAT
from gevent import pywsgi as wsgi
from werkzeug._reloader import run_with_reloader
from yaml import safe_load

from rowboat import ENV
from rowboat.sql import init_db
from rowboat.web import rowboat


class BotSupervisor(object):
    def __init__(self, env={}):
        self.proc = None
        self.env = env
        self.bind_signals()
        self.start()

    def bind_signals(self):
        signal.signal(signal.SIGUSR1, self.handle_sigusr1)

    def handle_sigusr1(self, signum, frame):
        print('SIGUSR1 - RESTARTING')
        gevent.spawn(self.restart)

    def start(self):
        env = copy.deepcopy(os.environ)
        env.update(self.env)
        self.proc = subprocess.Popen(['python', '-m', 'disco.cli', '--config', 'config.yaml'], env=env)

    def stop(self):
        self.proc.terminate()

    def restart(self):
        try:
            self.stop()
        except:
            pass

        self.start()

    def run_forever(self):
        while True:
            self.proc.wait()
            gevent.sleep(5)


@click.group()
def cli():
    logging.getLogger().setLevel(logging.INFO)


@cli.command()
@click.option('--reloader/--no-reloader', '-r', default=False)
def serve(reloader):
    def run():
        wsgi.WSGIServer(('0.0.0.0', 8686), rowboat).serve_forever()

    if reloader:
        run_with_reloader(run)
    else:
        run()


@cli.command()
@click.option('--env', '-e', default='local')
def bot(env):
    with open('config.yaml', 'r') as f:
        config = safe_load(f)

    supervisor = BotSupervisor(env={
        'ENV': env,
        'DSN': config['DSN'],
    })
    supervisor.run_forever()


@cli.command()
@click.option('--worker-id', '-w', default=0)
def workers(worker_id):
    from rowboat.tasks import TaskWorker

    # Log things to file
    file_handler = logging.FileHandler('worker-%s.log' % worker_id)
    log = logging.getLogger()
    file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
    log.addHandler(file_handler)

    for logname in ['peewee', 'requests']:
        logging.getLogger(logname).setLevel(logging.INFO)

    init_db(ENV)
    TaskWorker().run()


@cli.command('add-global-admin')
@click.argument('user-id')
def add_global_admin(user_id):
    from rowboat.models.user import User
    from rowboat.redis import rdb
    init_db(ENV)
    rdb.sadd('global_admins', user_id)
    user = User.get_id(user_id)
    user.update(admin=True)
    print('Ok, added {} as a global admin'.format(str(user))) # Not sure if I need to str here


@cli.command('wh-add')
@click.argument('guild-id')
@click.argument('flag')
def add_whitelist(guild_id, flag):
    from rowboat.models.guild import Guild
    init_db(ENV)

    flag = Guild.WhitelistFlags.get(flag)
    if not flag:
        print('Flag not found')
        return

    try:
        guild = Guild.get(guild_id=guild_id)
    except Guild.DoesNotExist:
        print('No guild exists with that id')
        return

    guild.whitelist.append(int(flag))
    guild.save()
    guild.emit_update()
    print('Added flag')


@cli.command('wh-rmv')
@click.argument('guild-id')
@click.argument('flag')
def rmv_whitelist(guild_id, flag):
    from rowboat.models.guild import Guild
    init_db(ENV)

    flag = Guild.WhitelistFlags.get(flag)
    if not flag:
        print('Flag not found')
        return

    try:
        guild = Guild.get(guild_id=guild_id)
    except Guild.DoesNotExist:
        print('No guild exists with that id')
        return

    guild.whitelist.remove(int(flag))
    guild.save()
    guild.emit_update()
    print('Removed flag')


if __name__ == '__main__':
    cli()