script.py 4.53 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
9
from typing import Any
Lukáš Lalinský's avatar
Lukáš Lalinský committed
10
from redis import Redis
11
from redis.sentinel import Sentinel as RedisSentinel
12
from optparse import OptionParser
Lukáš Lalinský's avatar
Lukáš Lalinský committed
13
from acoustid.config import Config
14
from acoustid.indexclient import IndexClientPool
15
from acoustid.utils import LocalSysLogHandler
16
from acoustid.db import DatabaseContext
Lukáš Lalinský's avatar
Lukáš Lalinský committed
17
from acoustid._release import GIT_RELEASE
Lukáš Lalinský's avatar
Lukáš Lalinský committed
18 19 20 21

logger = logging.getLogger(__name__)


22 23
class ScriptContext(object):

24 25 26
    def __init__(self, config, db, redis, index):
        # type: (Config, DatabaseContext, Redis, IndexClientPool) -> None
        self.config = config
27 28 29 30 31
        self.db = db
        self.redis = redis
        self.index = index

    def __enter__(self):
32
        # type: () -> ScriptContext
33 34 35
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
36
        # type: (Any, Any, Any) -> None
37 38 39
        self.db.close()


Lukáš Lalinský's avatar
Lukáš Lalinský committed
40 41
class Script(object):

Lukáš Lalinský's avatar
Lukáš Lalinský committed
42
    def __init__(self, config_path, tests=False):
Lukáš Lalinský's avatar
Lukáš Lalinský committed
43
        # type: (str, bool) -> None
44 45 46
        self.config = Config()
        if config_path:
            self.config.read(config_path)
47
        self.config.read_env(tests=tests)
48

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

Lukáš Lalinský's avatar
Lukáš Lalinský committed
52 53 54 55
        self.index = IndexClientPool(host=self.config.index.host,
                                     port=self.config.index.port,
                                     recycle=60)

56 57 58 59 60 61
        if self.config.redis.sentinel:
            self.redis_sentinel = RedisSentinel([(self.config.redis.host, self.config.redis.port)])
            self.redis = self.redis_sentinel.master_for(self.config.redis.cluster)  # type: Redis
        else:
            self.redis = Redis(host=self.config.redis.host,
                               port=self.config.redis.port)
62

63
        self._console_logging_configured = False
64 65
        if not tests:
            self.setup_logging()
66 67

    def setup_logging(self):
Lukáš Lalinský's avatar
Lukáš Lalinský committed
68
        # type: () -> None
69 70
        for logger_name, level in sorted(self.config.logging.levels.items()):
            logging.getLogger(logger_name).setLevel(level)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
71
        if self.config.logging.syslog:
72 73 74
            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
75
            logging.getLogger().addHandler(handler)
76 77
        else:
            self.setup_console_logging()
78

79
    def setup_console_logging(self, quiet=False):
Lukáš Lalinský's avatar
Lukáš Lalinský committed
80
        # type: (bool) -> None
81 82
        if self._console_logging_configured:
            return
83
        handler = logging.StreamHandler()
Lukáš Lalinský's avatar
Lukáš Lalinský committed
84
        handler.setFormatter(logging.Formatter('[%(asctime)s] [%(process)s] [%(levelname)s] %(message)s', '%Y-%m-%d %H:%M:%S %z'))
85 86
        if quiet:
            handler.setLevel(logging.ERROR)
87
        logging.getLogger().addHandler(handler)
88
        self._console_logging_configured = True
89

Lukáš Lalinský's avatar
Lukáš Lalinský committed
90
    def setup_sentry(self):
Lukáš Lalinský's avatar
Lukáš Lalinský committed
91
        # type: () -> None
Lukáš Lalinský's avatar
Lukáš Lalinský committed
92
        sentry_sdk.init(self.config.sentry.script_dsn, release=GIT_RELEASE)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
93

94
    def context(self):
95
        # type: () -> ScriptContext
96
        return ScriptContext(config=self.config, db=DatabaseContext(self), redis=self.redis, index=self.index)
97

98

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

125
    for engine in script.db_engines.values():
126 127
        engine.dispose()

128
    script.index.dispose()