| # pylint: disable-msg=C0111 |
| |
| import re, time, traceback |
| import common |
| from autotest_lib.client.common_lib import global_config |
| |
| RECONNECT_FOREVER = object() |
| |
| _DB_EXCEPTIONS = ('DatabaseError', 'OperationalError', 'ProgrammingError') |
| _GLOBAL_CONFIG_NAMES = { |
| 'username' : 'user', |
| 'db_name' : 'database', |
| } |
| |
| def _copy_exceptions(source, destination): |
| for exception_name in _DB_EXCEPTIONS: |
| try: |
| setattr(destination, exception_name, |
| getattr(source, exception_name)) |
| except AttributeError: |
| # Under the django backend: |
| # Django 1.3 does not have OperationalError and ProgrammingError. |
| # Let's just mock these classes with the base DatabaseError. |
| setattr(destination, exception_name, |
| getattr(source, 'DatabaseError')) |
| |
| |
| class _GenericBackend(object): |
| def __init__(self, database_module): |
| self._database_module = database_module |
| self._connection = None |
| self._cursor = None |
| self.rowcount = None |
| _copy_exceptions(database_module, self) |
| |
| |
| def connect(self, host=None, username=None, password=None, db_name=None): |
| """ |
| This is assumed to enable autocommit. |
| """ |
| raise NotImplementedError |
| |
| |
| def disconnect(self): |
| if self._connection: |
| self._connection.close() |
| self._connection = None |
| self._cursor = None |
| |
| |
| def execute(self, query, parameters=None): |
| if parameters is None: |
| parameters = () |
| self._cursor.execute(query, parameters) |
| self.rowcount = self._cursor.rowcount |
| return self._cursor.fetchall() |
| |
| |
| class _MySqlBackend(_GenericBackend): |
| def __init__(self): |
| import MySQLdb |
| super(_MySqlBackend, self).__init__(MySQLdb) |
| |
| |
| @staticmethod |
| def convert_boolean(boolean, conversion_dict): |
| 'Convert booleans to integer strings' |
| return str(int(boolean)) |
| |
| |
| def connect(self, host=None, username=None, password=None, db_name=None): |
| import MySQLdb.converters |
| convert_dict = MySQLdb.converters.conversions |
| convert_dict.setdefault(bool, self.convert_boolean) |
| |
| self._connection = self._database_module.connect( |
| host=host, user=username, passwd=password, db=db_name, |
| conv=convert_dict) |
| self._connection.autocommit(True) |
| self._cursor = self._connection.cursor() |
| |
| |
| class _SqliteBackend(_GenericBackend): |
| def __init__(self): |
| try: |
| from pysqlite2 import dbapi2 |
| except ImportError: |
| from sqlite3 import dbapi2 |
| super(_SqliteBackend, self).__init__(dbapi2) |
| self._last_insert_id_re = re.compile(r'\sLAST_INSERT_ID\(\)', |
| re.IGNORECASE) |
| |
| |
| def connect(self, host=None, username=None, password=None, db_name=None): |
| self._connection = self._database_module.connect(db_name) |
| self._connection.isolation_level = None # enable autocommit |
| self._cursor = self._connection.cursor() |
| |
| |
| def execute(self, query, parameters=None): |
| # pysqlite2 uses paramstyle=qmark |
| # TODO: make this more sophisticated if necessary |
| query = query.replace('%s', '?') |
| # pysqlite2 can't handle parameters=None (it throws a nonsense |
| # exception) |
| if parameters is None: |
| parameters = () |
| # sqlite3 doesn't support MySQL's LAST_INSERT_ID(). Instead it has |
| # something similar called LAST_INSERT_ROWID() that will do enough of |
| # what we want (for our non-concurrent unittest use case). |
| query = self._last_insert_id_re.sub(' LAST_INSERT_ROWID()', query) |
| return super(_SqliteBackend, self).execute(query, parameters) |
| |
| |
| class _DjangoBackend(_GenericBackend): |
| def __init__(self): |
| from django.db import backend, connection, transaction |
| import django.db as django_db |
| super(_DjangoBackend, self).__init__(django_db) |
| self._django_connection = connection |
| self._django_transaction = transaction |
| |
| |
| def connect(self, host=None, username=None, password=None, db_name=None): |
| self._connection = self._django_connection |
| self._cursor = self._connection.cursor() |
| |
| |
| def execute(self, query, parameters=None): |
| try: |
| return super(_DjangoBackend, self).execute(query, |
| parameters=parameters) |
| finally: |
| self._django_transaction.commit_unless_managed() |
| |
| |
| _BACKEND_MAP = { |
| 'mysql': _MySqlBackend, |
| 'sqlite': _SqliteBackend, |
| 'django': _DjangoBackend, |
| } |
| |
| |
| class DatabaseConnection(object): |
| """ |
| Generic wrapper for a database connection. Supports both mysql and sqlite |
| backends. |
| |
| Public attributes: |
| * reconnect_enabled: if True, when an OperationalError occurs the class will |
| try to reconnect to the database automatically. |
| * reconnect_delay_sec: seconds to wait before reconnecting |
| * max_reconnect_attempts: maximum number of time to try reconnecting before |
| giving up. Setting to RECONNECT_FOREVER removes the limit. |
| * rowcount - will hold cursor.rowcount after each call to execute(). |
| * global_config_section - the section in which to find DB information. this |
| should be passed to the constructor, not set later, and may be None, in |
| which case information must be passed to connect(). |
| * debug - if set True, all queries will be printed before being executed |
| """ |
| _DATABASE_ATTRIBUTES = ('db_type', 'host', 'username', 'password', |
| 'db_name') |
| |
| def __init__(self, global_config_section=None, debug=False): |
| self.global_config_section = global_config_section |
| self._backend = None |
| self.rowcount = None |
| self.debug = debug |
| |
| # reconnect defaults |
| self.reconnect_enabled = True |
| self.reconnect_delay_sec = 20 |
| self.max_reconnect_attempts = 10 |
| |
| self._read_options() |
| |
| |
| def _get_option(self, name, provided_value, use_afe_setting=False): |
| """Get value of given option from global config. |
| |
| @param name: Name of the config. |
| @param provided_value: Value being provided to override the one from |
| global config. |
| @param use_afe_setting: Force to use the settings in AFE, default is |
| False. |
| """ |
| # TODO(dshi): This function returns the option value depends on multiple |
| # conditions. The value of `provided_value` has highest priority, then |
| # the code checks if use_afe_setting is True, if that's the case, force |
| # to use settings in AUTOTEST_WEB. At last the value is retrieved from |
| # specified global config section. |
| # The logic is too complicated for a generic function named like |
| # _get_option. Ideally we want to make it clear from caller that it |
| # wants to get database credential from one of the 3 ways: |
| # 1. Use the credential from given config section |
| # 2. Use the credential from AUTOTEST_WEB section |
| # 3. Use the credential provided by caller. |
| if provided_value is not None: |
| return provided_value |
| section = ('AUTOTEST_WEB' if use_afe_setting else |
| self.global_config_section) |
| if section: |
| global_config_name = _GLOBAL_CONFIG_NAMES.get(name, name) |
| return global_config.global_config.get_config_value( |
| section, global_config_name) |
| |
| return getattr(self, name, None) |
| |
| |
| def _read_options(self, db_type=None, host=None, username=None, |
| password=None, db_name=None): |
| """Read database information from global config. |
| |
| Unless any parameter is specified a value, the connection will use |
| database name from given configure section (self.global_config_section), |
| and database credential from AFE database settings (AUTOTEST_WEB). |
| |
| @param db_type: database type, default to None. |
| @param host: database hostname, default to None. |
| @param username: user name for database connection, default to None. |
| @param password: database password, default to None. |
| @param db_name: database name, default to None. |
| """ |
| self.db_name = self._get_option('db_name', db_name) |
| use_afe_setting = not bool(db_type or host or username or password) |
| |
| # Database credential can be provided by the caller, as passed in from |
| # function connect. |
| self.db_type = self._get_option('db_type', db_type, use_afe_setting) |
| self.host = self._get_option('host', host, use_afe_setting) |
| self.username = self._get_option('username', username, use_afe_setting) |
| self.password = self._get_option('password', password, use_afe_setting) |
| |
| |
| def _get_backend(self, db_type): |
| if db_type not in _BACKEND_MAP: |
| raise ValueError('Invalid database type: %s, should be one of %s' % |
| (db_type, ', '.join(_BACKEND_MAP.keys()))) |
| backend_class = _BACKEND_MAP[db_type] |
| return backend_class() |
| |
| |
| def _reached_max_attempts(self, num_attempts): |
| return (self.max_reconnect_attempts is not RECONNECT_FOREVER and |
| num_attempts > self.max_reconnect_attempts) |
| |
| |
| def _is_reconnect_enabled(self, supplied_param): |
| if supplied_param is not None: |
| return supplied_param |
| return self.reconnect_enabled |
| |
| |
| def _connect_backend(self, try_reconnecting=None): |
| num_attempts = 0 |
| while True: |
| try: |
| self._backend.connect(host=self.host, username=self.username, |
| password=self.password, |
| db_name=self.db_name) |
| return |
| except self._backend.OperationalError: |
| num_attempts += 1 |
| if not self._is_reconnect_enabled(try_reconnecting): |
| raise |
| if self._reached_max_attempts(num_attempts): |
| raise |
| traceback.print_exc() |
| print ("Can't connect to database; reconnecting in %s sec" % |
| self.reconnect_delay_sec) |
| time.sleep(self.reconnect_delay_sec) |
| self.disconnect() |
| |
| |
| def connect(self, db_type=None, host=None, username=None, password=None, |
| db_name=None, try_reconnecting=None): |
| """ |
| Parameters passed to this function will override defaults from global |
| config. try_reconnecting, if passed, will override |
| self.reconnect_enabled. |
| """ |
| self.disconnect() |
| self._read_options(db_type, host, username, password, db_name) |
| |
| self._backend = self._get_backend(self.db_type) |
| _copy_exceptions(self._backend, self) |
| self._connect_backend(try_reconnecting) |
| |
| |
| def disconnect(self): |
| if self._backend: |
| self._backend.disconnect() |
| |
| |
| def execute(self, query, parameters=None, try_reconnecting=None): |
| """ |
| Execute a query and return cursor.fetchall(). try_reconnecting, if |
| passed, will override self.reconnect_enabled. |
| """ |
| if self.debug: |
| print 'Executing %s, %s' % (query, parameters) |
| # _connect_backend() contains a retry loop, so don't loop here |
| try: |
| results = self._backend.execute(query, parameters) |
| except self._backend.OperationalError: |
| if not self._is_reconnect_enabled(try_reconnecting): |
| raise |
| traceback.print_exc() |
| print ("MYSQL connection died; reconnecting") |
| self.disconnect() |
| self._connect_backend(try_reconnecting) |
| results = self._backend.execute(query, parameters) |
| |
| self.rowcount = self._backend.rowcount |
| return results |
| |
| |
| def get_database_info(self): |
| return dict((attribute, getattr(self, attribute)) |
| for attribute in self._DATABASE_ATTRIBUTES) |
| |
| |
| @classmethod |
| def get_test_database(cls, file_path=':memory:', **constructor_kwargs): |
| """ |
| Factory method returning a DatabaseConnection for a temporary in-memory |
| database. |
| """ |
| database = cls(**constructor_kwargs) |
| database.reconnect_enabled = False |
| database.connect(db_type='sqlite', db_name=file_path) |
| return database |
| |
| |
| class TranslatingDatabase(DatabaseConnection): |
| """ |
| Database wrapper than applies arbitrary substitution regexps to each query |
| string. Useful for SQLite testing. |
| """ |
| def __init__(self, translators): |
| """ |
| @param translation_regexps: list of callables to apply to each query |
| string (in order). Each accepts a query string and returns a |
| (possibly) modified query string. |
| """ |
| super(TranslatingDatabase, self).__init__() |
| self._translators = translators |
| |
| |
| def execute(self, query, parameters=None, try_reconnecting=None): |
| for translator in self._translators: |
| query = translator(query) |
| return super(TranslatingDatabase, self).execute( |
| query, parameters=parameters, try_reconnecting=try_reconnecting) |
| |
| |
| @classmethod |
| def make_regexp_translator(cls, search_re, replace_str): |
| """ |
| Returns a translator that calls re.sub() on the query with the given |
| search and replace arguments. |
| """ |
| def translator(query): |
| return re.sub(search_re, replace_str, query) |
| return translator |
