class Command(object):
    def __init__(self, options):
        self.options = options
        self.ensure_bootstrapped()
        self.run()

    @staticmethod
    def configure(parser, subparsers):
        raise NotImplemented('Every command needs a static configure(...) method')

    def ensure_bootstrapped(self):
        if not bootstrap.check_bootstrapped():
            bootstrap.bootstrap()

    def run():
        raise NotImplemented('Every command needs a run method, otherwise not much is going to happen') def setup(log_path=None, stdout_enabled=True):
def update_verbosity(debug=False, extra_verbose=False):
    def __call__(self, parser, namespace, value, option_string=None):
    def __call__(self, parser, namespace, value, option_string=None):
    def __call__(self, parser, namespace, value, option_string=None):
def configure(parser):
def get_logger(name, level=None):
    def format(self, record):
    def __init__(self, name, level=logging.DEBUG):
class _SetDebug(argparse.Action):
class _SetStdout(argparse.Action):
class _SetLogPath(argparse.Action):
class Formatter(logging.Formatter):
class Logger(logging.Logger): class InvalidDatabaseException(Exception):
class AWSBucketNameInvalid(Exception):
class KeyFilterEmpty(Exception):
class MissingTable(Exception):
RE_VALIDATE_BUCKET_NAME = re.compile(r"^(?:[a-z0-9]+(?:[.\-]?[a-z0-9]+)*)+$")
def prompt(message, default=None, validate=None, string_escape=False):
def confirm(message, default=True):
    def _check(value): class WorkerPool(object):
    def __init__(self, num_workers=None, idle_timeout=None):
        self.logger = log.get_logger('WorkerPool')
        self.num_workers = num_workers or max(1, int(multiprocessing.cpu_count() * 0.8))
        self.idle_timeout = idle_timeout
        self._workers = []
        self.pid = os.getpid()
        self._worker_lock = multiprocessing.Lock()
        self._last_work_time = time.time()

    def poll(self):
        running = [worker for worker in self._workers if worker.is_alive()]

        if self.idle_timeout is not None:
            if any([worker.is_working() for worker in self._workers]):
                self._last_work_time = time.time()
            elif time.time() > (self._last_work_time + self.idle_timeout):
                return False

        diff = self.num_workers - len(running)
        if diff > 0:
            self.logger.debug('Starting %d workers, for a total of %d', diff, self.num_workers)
            with LoaderStorage.fork_wrapper():
                running += [self._start_worker(i) for i in xrange(diff)]
        self._workers = running

        return True

    def stop(self):
        [worker.signal_exit() for worker in self._workers if worker.is_alive()]
        [worker.join() for worker in self._workers if worker.is_alive()]

    def _start_worker(self, index):
        worker = Worker(index * 0.1, self.pid, self._worker_lock)
        worker.start()
        return worker _POOL = connection_pool.ConnectionPool()
    _POOL = connection_pool.ConnectionPool()
def recreate_pool():
def close_connections():
def get_connection(host, port, database, user, password, pooled=True, **kwargs): MEMSQL_LOADER_DB = 'memsql_loader.db'
def get_loader_db_path():
    def __new__(cls, *args, **kwargs):
    def drop_database(cls):
    def fork_wrapper(cls):
    def __init__(self):
class LoaderStorage(APSWStorage): def daemonize(logfile):
    """ Daemonize the current process.
    When it returns you will be running in a daemonized child process.
    Mostly copied from: http://www.jejik.com/articles/2007/02/a_simple_unix_linux_daemon_in_python/
    """

    try:
        pid = os.fork()
        if pid > 0:
            sys.exit(0)
    except OSError as e:
        sys.stderr.write("Fork #1 failed: %d (%s)\n" % (e.errno, e.strerror))
        sys.exit(1)

    os.chdir("/")
    os.setsid()
    os.umask(0)

    try:
        pid = os.fork()
        if pid > 0:
            sys.exit(0)
    except OSError as e:
        sys.stderr.write("Fork #2 failed: %d (%s)\n" % (e.errno, e.strerror))
        sys.exit(1)

    sys.stdout.flush()
    sys.stderr.flush()

    # setup a new STDIN
    si = file('/dev/null', 'r')
    os.dup2(si.fileno(), sys.stdin.fileno())

    # setup new STDOUT
    output = '/dev/null' if logfile is None else logfile
    so = file(output, 'a+')
    se = file(output, 'a+', 0)
    os.dup2(so.fileno(), sys.stdout.fileno())
    os.dup2(se.fileno(), sys.stderr.fileno()) def setuser(user_or_pid):
    try:
        try:
            userinfo = pwd.getpwuid(int(user_or_pid))
        except ValueError:
            userinfo = pwd.getpwnam(user_or_pid)

        __fix_perms(userinfo)
        os.setgid(userinfo.pw_gid)
        os.setuid(userinfo.pw_uid)
        return True
    except KeyError as e:
        print(str(e))
        return False def bootstrap(force=False):
    logger = log.get_logger('Bootstrap')  # noqa

    def write_log(title, name, msg):
        log_title_width = 28
        title = ("%s [%s]: " % (title, name)).rjust(log_title_width, ' ')
        logger.info(title + msg)

    write_log('Database', storage.MEMSQL_LOADER_DB, 'Checking...')
    if force:
        write_log('Database', storage.MEMSQL_LOADER_DB, 'Dropping...')
        storage.LoaderStorage.drop_database()
    write_log('Database', storage.MEMSQL_LOADER_DB, 'Ready.')

    for Model in MODELS.values():
        instance = Model()
        if not instance.ready():
            write_log('Table', Model.__name__, 'Bootstrapping...')
            instance.setup()
        write_log('Table', Model.__name__, 'Ready.') def get_pid_file_path():
def delete_pid_file():
def write_pid_file():
def get_server_pid():
def is_server_running():
from memsql_loader.util.command import Command
from memsql_loader.util import log, cli_utils
from memsql_loader.execution.worker_pool import WorkerPool
from memsql_loader.db import pool
from memsql_loader.loader_db import storage
from memsql_loader.util.daemonize import daemonize
from memsql_loader.util.setuser import setuser
from memsql_loader.util import bootstrap, servers
import argparse
import multiprocessing
import time
import sys
import signal
""" Start the worker server """


WORKER_WARN_THRESHOLD = 100

# This class is used in the load command to start a server with default
# arguments in a separate process.
class ServerProcess(multiprocessing.Process):
    def __init__(self, daemonize=False, num_workers=None, idle_timeout=None):
        self.num_workers = num_workers
        self.idle_timeout = idle_timeout
        super(ServerProcess, self).__init__()
        self.daemonize = daemonize

    def run(self):
        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(parser_class=argparse.ArgumentParser)
        Server.configure(parser, subparsers)
        log.configure(parser)
        fake_args = ['server']
        if self.num_workers is not None:
            fake_args.append('--num-workers')
            fake_args.append(str(self.num_workers))
        if self.idle_timeout is not None:
            fake_args.append('--idle-timeout')
            fake_args.append(str(self.idle_timeout))
        options = parser.parse_args(fake_args)
        options.daemonize = self.daemonize
        Server(options)

class Server(Command):
    @staticmethod
    def configure(parser, subparsers):
        subparser = subparsers.add_parser('server', help='Start a server that runs queued jobs')
        subparser.set_defaults(command=Server)
        subparser.add_argument('-d', '--daemonize', action='store_true', help='Daemonize the MemSQL Loader server process.', default=False)
        subparser.add_argument('--set-user', default=None, help='Specify a user for MemSQL Loader to use.')
        subparser.add_argument('-n', '--num-workers', default=None, type=int,
            help='Number of workers to run; equates to the number of loads that can be run in parallel.')
        subparser.add_argument('-i', '--idle-timeout', default=None, type=int,
            help='Seconds before server automatically shuts down; defaults to never.')
        subparser.add_argument('-f', '--force-workers', action='store_true',
            help='Ignore warnings on number of workers. This is potentially dangerous!')

    def ensure_bootstrapped(self):
        if not bootstrap.check_bootstrapped():
            bootstrap.bootstrap()

    def run(self):
        signal.signal(signal.SIGINT, self.stop)
        signal.signal(signal.SIGQUIT, self.stop)
        signal.signal(signal.SIGTERM, self.stop)

        self.exiting = False

        self.logger = log.get_logger('Server')