"""Connector to Mysql database.""" import contextlib import os import types from sqlalchemy import create_engine from sqlalchemy import event from sqlalchemy import exc from sqlalchemy import select from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker from sqlalchemy.pool import QueuePool, StaticPool DEFAULT_DB_NAME = 'default' default = None # type: MySQLConnection _db_config = None _environment = None connections = types.SimpleNamespace() _url_template = ( 'mysql+pymysql://{user}:{password}@{host}/{db}?charset={charset}') CONFIG_DEFAULTS = { 'charset': 'utf8', 'encoding': 'utf-8', 'pool_class': QueuePool, 'pool_pre_ping': True, 'pool_size': 5, 'max_overflow': -1, 'pool_recycle': 3600, } class ImproperlyConfigured(Exception): """Exception for invalid environment configuration.""" class MySQLConnection: """MySQL connection class.""" def __init__(self, name='default', base_model=None): """Init new DB connection with specified name.""" if _db_config is None: raise ImproperlyConfigured( 'Please call db.configure() before creating connections') self.name = name if name not in _db_config: raise ImproperlyConfigured( 'Connection {} name is not configured'.format(name)) self.config = CONFIG_DEFAULTS.copy() self.config.update(_db_config[name] or {}) self.engine = self._create_engine() self.session_maker = sessionmaker(bind=self.engine) if base_model: self.BaseModel = base_model self.BaseModel.metadata.bind = self.engine else: self.BaseModel = declarative_base(bind=self.engine) setattr(connections, name, self) def _create_engine(self): """Get DB connection URL for ORM.""" if 'url' in self.config: url = self.config.get('url') else: url = _url_template.format( user=self.config.get('user'), password=self.config.get('password'), host=self.config.get('host'), db=self.config.get('database'), charset=self.config.get('charset')) pool_class = self.config.get('pool_class') engine = create_engine( url, connect_args=self.config.get('connect_args') or {}, poolclass=pool_class, pool_size=self.config.get('pool_size'), max_overflow=self.config.get('max_overflow'), pool_recycle=self.config.get('pool_recycle'), encoding=self.config.get('encoding')) if self.config.get('pool_pre_ping'): event.listen(engine, 'engine_connect', self._ping_connection) return engine def _ping_connection(self, connection, branch): """Ping database connection after engine_connect event. This function is copied verbatim from http://docs.sqlalchemy.org/en/latest/core/pooling.html """ if branch: # pragma: no cover # "branch" refers to a sub-connection of a connection, # we don't want to bother pinging on these. return # turn off "close with result". This flag is only used with # "connectionless" execution, otherwise will be False in any case save_should_close_with_result = connection.should_close_with_result connection.should_close_with_result = False try: # run a SELECT 1. use a core select() so that # the SELECT of a scalar value without a table is # appropriately formatted for the backend connection.scalar(select([1])) except exc.DBAPIError as err: # catch SQLAlchemy's DBAPIError, which is a wrapper # for the DBAPI's exception. It includes a .connection_invalidated # attribute which specifies if this connection is a "disconnect" # condition, which is based on inspection of the original exception # by the dialect in use. if err.connection_invalidated: # run the same SELECT again - the connection will re-validate # itself and establish a new connection. The disconnect # detection here also causes the whole connection pool to be # invalidated so that all stale connections are discarded. connection.scalar(select([1])) else: raise finally: # restore "close with result" connection.should_close_with_result = save_should_close_with_result @contextlib.contextmanager def session(self, read_only=False): """Create DB session for this connection.""" session_obj = self.session_maker() try: yield session_obj # Only call commit if session contains a write transaction. if not read_only: session_obj.commit() except: session_obj.rollback() raise finally: session_obj.close() def create_all(self): """Create all known tables definitions for this connection.""" raise ImproperlyConfigured( 'Schema creation is only for test environment') def drop_all(self): """Drop all known tables definitions for this connection.""" raise ImproperlyConfigured( 'Dropping schema is only for test environment') class TestConnection(MySQLConnection): """Connection class for Test environment.""" def _create_engine(self): """Get DB connection URL for ORM.""" config = _db_config.get('test') if config is None and 'TEST_DB_HOST' in os.environ: config = { 'user': os.environ.get('TEST_DB_USER'), 'password': os.environ.get('TEST_DB_PASSWORD'), 'host': os.environ.get('TEST_DB_HOST'), 'db': os.environ.get('TEST_DB_NAME'), } if config is not None: url = _url_template.format(charset='utf8', **config) else: url = 'sqlite://' pool_class = self.config.get('pool_class') or StaticPool return create_engine( url, connect_args=self.config.get('connect_args') or {}, poolclass=pool_class) def create_all(self): """Create all known tables definitions for this connection.""" self.BaseModel.metadata.create_all(self.engine) def drop_all(self): """Drop all known tables definitions for this connection.""" self.BaseModel.metadata.drop_all(self.engine) def configure(environment, config=None, base_model=None): """Configure DB layer. Args: environment (str): Environment to run, e.g test/dev config (dict): DB configuration. base_model (declarative_base): Base model class (optional). """ global _db_config global _environment global default if not config: config = {'default': None} _db_config = config _environment = environment if _environment == 'test': connection_cls = TestConnection else: connection_cls = MySQLConnection if DEFAULT_DB_NAME in config: default = connection_cls('default', base_model=base_model)