1#pylint: disable-msg=C0111
2
3# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Scheduler helper libraries.
8"""
9import logging
10import os
11
12import common
13
14from autotest_lib.client.common_lib import global_config
15from autotest_lib.client.common_lib import logging_config
16from autotest_lib.client.common_lib import logging_manager
17from autotest_lib.client.common_lib import utils
18from autotest_lib.database import database_connection
19from autotest_lib.frontend import setup_django_environment
20from autotest_lib.frontend.afe import readonly_connection
21from autotest_lib.server import utils as server_utils
22
23
24DB_CONFIG_SECTION = 'AUTOTEST_WEB'
25
26# Translations necessary for scheduler queries to work with SQLite.
27# Though this is only used for testing it is included in this module to avoid
28# circular imports.
29_re_translator = database_connection.TranslatingDatabase.make_regexp_translator
30_DB_TRANSLATORS = (
31        _re_translator(r'NOW\(\)', 'time("now")'),
32        _re_translator(r'LAST_INSERT_ID\(\)', 'LAST_INSERT_ROWID()'),
33        # older SQLite doesn't support group_concat, so just don't bother until
34        # it arises in an important query
35        _re_translator(r'GROUP_CONCAT\((.*?)\)', r'\1'),
36        _re_translator(r'TRUNCATE TABLE', 'DELETE FROM'),
37        _re_translator(r'ISNULL\(([a-z,_]+)\)',
38                       r'ifnull(nullif(\1, NULL), \1) DESC'),
39)
40
41
42class SchedulerError(Exception):
43    """Raised by the scheduler when an inconsistent state occurs."""
44
45
46class ConnectionManager(object):
47    """Manager for the django database connections.
48
49    The connection is used through scheduler_models and monitor_db.
50    """
51    __metaclass__ = server_utils.Singleton
52
53    def __init__(self, readonly=True, autocommit=True):
54        """Set global django database options for correct connection handling.
55
56        @param readonly: Globally disable readonly connections.
57        @param autocommit: Initialize django autocommit options.
58        """
59        self.db_connection = None
60        # bypass the readonly connection
61        readonly_connection.set_globally_disabled(readonly)
62        if autocommit:
63            # ensure Django connection is in autocommit
64            setup_django_environment.enable_autocommit()
65
66
67    @classmethod
68    def open_connection(cls):
69        """Open a new database connection.
70
71        @return: An instance of the newly opened connection.
72        """
73        db = database_connection.DatabaseConnection(DB_CONFIG_SECTION)
74        db.connect(db_type='django')
75        return db
76
77
78    def get_connection(self):
79        """Get a connection.
80
81        @return: A database connection.
82        """
83        if self.db_connection is None:
84            self.db_connection = self.open_connection()
85        return self.db_connection
86
87
88    def disconnect(self):
89        """Close the database connection."""
90        try:
91            self.db_connection.disconnect()
92        except Exception as e:
93            logging.debug('Could not close the db connection. %s', e)
94
95
96    def __del__(self):
97        self.disconnect()
98
99
100class SchedulerLoggingConfig(logging_config.LoggingConfig):
101    """Configure timestamped logging for a scheduler."""
102    GLOBAL_LEVEL = logging.INFO
103
104    @classmethod
105    def get_log_name(cls, timestamped_logfile_prefix):
106        """Get the name of a logfile.
107
108        @param timestamped_logfile_prefix: The prefix to apply to the
109            a timestamped log. Eg: 'scheduler' will create a logfile named
110            scheduler.log.2014-05-12-17.24.02.
111
112        @return: The timestamped log name.
113        """
114        return cls.get_timestamped_log_name(timestamped_logfile_prefix)
115
116
117    def configure_logging(self, log_dir=None, logfile_name=None,
118                          timestamped_logfile_prefix='scheduler'):
119        """Configure logging to a specified logfile.
120
121        @param log_dir: The directory to log into.
122        @param logfile_name: The name of the log file.
123        @timestamped_logfile_prefix: The prefix to apply to the name of
124            the logfile, if a log file name isn't specified.
125        """
126        super(SchedulerLoggingConfig, self).configure_logging(use_console=True)
127
128        if log_dir is None:
129            log_dir = self.get_server_log_dir()
130        if not logfile_name:
131            logfile_name = self.get_log_name(timestamped_logfile_prefix)
132
133        self.add_file_handler(logfile_name, logging.DEBUG, log_dir=log_dir)
134        symlink_path = os.path.join(
135                log_dir, '%s.latest' % timestamped_logfile_prefix)
136        try:
137            os.unlink(symlink_path)
138        except OSError:
139            pass
140        os.symlink(os.path.join(log_dir, logfile_name), symlink_path)
141
142
143def setup_logging(log_dir, log_name, timestamped_logfile_prefix='scheduler'):
144    """Setup logging to a given log directory and log file.
145
146    @param log_dir: The directory to log into.
147    @param log_name: Name of the log file.
148    @param timestamped_logfile_prefix: The prefix to apply to the logfile.
149    """
150    logging_manager.configure_logging(
151            SchedulerLoggingConfig(), log_dir=log_dir, logfile_name=log_name,
152            timestamped_logfile_prefix=timestamped_logfile_prefix)
153
154
155def check_production_settings(scheduler_options):
156    """Check the scheduler option's production settings.
157
158    @param scheduler_options: Settings for scheduler.
159
160    @raises SchedulerError: If a loclhost scheduler is started with
161       production settings.
162    """
163    db_server = global_config.global_config.get_config_value('AUTOTEST_WEB',
164                                                             'host')
165    if (not scheduler_options.production and
166        not utils.is_localhost(db_server)):
167        raise SchedulerError('Scheduler is not running in production mode, you '
168                             'should not set database to hosts other than '
169                             'localhost. It\'s currently set to %s.\nAdd option'
170                             ' --production if you want to skip this check.' %
171                             db_server)
172