script.py 4.04 KB
Newer Older
Lukáš Lalinský's avatar
Lukáš Lalinský committed
1
# Copyright (C) 2011 Lukas Lalinsky
2
# Distributed under the MIT license, see the LICENSE file for details.
Lukáš Lalinský's avatar
Lukáš Lalinský committed
3

4
import sys
Lukáš Lalinský's avatar
Lukáš Lalinský committed
5 6
import logging
import sqlalchemy
7
import sqlalchemy.pool
Lukáš Lalinský's avatar
Lukáš Lalinský committed
8
import sentry_sdk
Lukáš Lalinský's avatar
Lukáš Lalinský committed
9
from redis import Redis
10
from optparse import OptionParser
Lukáš Lalinský's avatar
Lukáš Lalinský committed
11
from acoustid.config import Config
12
from acoustid.indexclient import IndexClientPool
13
from acoustid.utils import LocalSysLogHandler
Lukáš Lalinský's avatar
Lukáš Lalinský committed
14
from acoustid.db import DatabaseContext
Lukáš Lalinský's avatar
Lukáš Lalinský committed
15
from acoustid._release import GIT_RELEASE
Lukáš Lalinský's avatar
Lukáš Lalinský committed
16 17 18 19

logger = logging.getLogger(__name__)


Lukáš Lalinský's avatar
Lukáš Lalinský committed
20 21 22 23 24 25 26 27 28 29 30 31 32 33
class ScriptContext(object):

    def __init__(self, db, redis, index):
        self.db = db
        self.redis = redis
        self.index = index

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.db.close()


Lukáš Lalinský's avatar
Lukáš Lalinský committed
34 35
class Script(object):

Lukáš Lalinský's avatar
Lukáš Lalinský committed
36
    def __init__(self, config_path, tests=False):
37 38 39
        self.config = Config()
        if config_path:
            self.config.read(config_path)
40
        self.config.read_env(tests=tests)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
41 42 43 44

        create_engine_kwargs = {'poolclass': sqlalchemy.pool.AssertionPool} if tests else {}
        self.db_engines = self.config.databases.create_engines(**create_engine_kwargs)

45 46 47 48 49
        if not self.config.index.host:
            self.index = None
        else:
            self.index = IndexClientPool(host=self.config.index.host,
                                         port=self.config.index.port,
50
                                         recycle=60)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
51

52 53 54 55
        if not self.config.redis.host:
            self.redis = None
        else:
            self.redis = Redis(host=self.config.redis.host,
56
                               port=self.config.redis.port)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
57

58
        self._console_logging_configured = False
59 60
        self.setup_logging()

Lukáš Lalinský's avatar
Lukáš Lalinský committed
61 62
    @property
    def engine(self):
63
        return self.db_engines['main']
Lukáš Lalinský's avatar
Lukáš Lalinský committed
64

65 66 67
    def setup_logging(self):
        for logger_name, level in sorted(self.config.logging.levels.items()):
            logging.getLogger(logger_name).setLevel(level)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
68
        if self.config.logging.syslog:
69 70 71
            handler = LocalSysLogHandler(ident='acoustid',
                facility=self.config.logging.syslog_facility, log_pid=True)
            handler.setFormatter(logging.Formatter('%(name)s: %(message)s'))
Lukáš Lalinský's avatar
Lukáš Lalinský committed
72
            logging.getLogger().addHandler(handler)
73 74
        else:
            self.setup_console_logging()
75

76
    def setup_console_logging(self, quiet=False):
77 78
        if self._console_logging_configured:
            return
79 80
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter('%(asctime)s [%(levelname)s] %(name)s - %(message)s', '%H:%M:%S'))
81 82
        if quiet:
            handler.setLevel(logging.ERROR)
83
        logging.getLogger().addHandler(handler)
84
        self._console_logging_configured = True
85

Lukáš Lalinský's avatar
Lukáš Lalinský committed
86
    def setup_sentry(self):
Lukáš Lalinský's avatar
Lukáš Lalinský committed
87
        sentry_sdk.init(self.config.sentry.script_dsn, release=GIT_RELEASE)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
88

Lukáš Lalinský's avatar
Lukáš Lalinský committed
89 90 91 92
    def context(self):
        db = DatabaseContext(self).session
        return ScriptContext(db=db, redis=self.redis, index=self.index)

93

94
def run_script(func, option_cb=None, master_only=False):
95 96 97
    parser = OptionParser()
    parser.add_option("-c", "--config", dest="config",
        help="configuration file", metavar="FILE")
98 99
    parser.add_option("-q", "--quiet", dest="quiet", action="store_true",
        default=False, help="don't print info messages to stdout")
100 101
    if option_cb is not None:
        option_cb(parser)
102 103 104 105
    (options, args) = parser.parse_args()
    if not options.config:
        parser.error('no configuration file')
    script = Script(options.config)
106
    script.setup_console_logging(options.quiet)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
107
    script.setup_sentry()
108 109
    if master_only and script.config.cluster.role != 'master':
        logger.debug("Not running script %s on a slave server", sys.argv[0])
110
    else:
111 112 113
        logger.debug("Running script %s", sys.argv[0])
        try:
            func(script, options, args)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
114
        except Exception:
115 116 117 118
            logger.exception("Script finished %s with an exception", sys.argv[0])
            raise
        else:
            logger.debug("Script finished %s successfuly", sys.argv[0])
119
    script.engine.dispose()
120
    script.index.dispose()