from contextlib import contextmanager

import sqlalchemy
from sqlalchemy import event
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base


def _include_sqlalchemy(obj):
    for module in sqlalchemy, sqlalchemy.orm:
        for key in module.__all__:
            if not hasattr(obj, key):
                setattr(obj, key, getattr(module, key))
    obj.relationship = obj.relationship
    obj.relation = obj.relation
    obj.dynamic_loader = obj.dynamic_loader
    obj.event = event


class _SQLAlchemy:
    def __init__(self, mysql_user, mysql_password, mysql_host, mysql_db):
        _include_sqlalchemy(self)
        engine = create_engine(
            'mysql+mysqldb://{user}:{password}@{host}:3306/{db}?charset=utf8'.format(
                user=mysql_user, password=mysql_password, host=mysql_host, db=mysql_db),
            encoding='utf8'
        )
        self.session = self.sessionmaker(bind=engine)()

    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise


_BaseModel = declarative_base()