from sqlalchemy import text

DB_ENGINE = None


def init_db_engine(db):
    global DB_ENGINE
    DB_ENGINE = db.engine


def db_conn():
    """
    일반 디비 쿼리 커넥션 생성
    사용방법
        with db_conn() as conn:
            dao.insert_any(conn, data)
    """
    global DB_ENGINE
    return ExConnection(DB_ENGINE)


def db_trans():
    """
    트랜잭션 디비 쿼리 커넥션 생성
    사용방법
        with db_trans() as conn:
            dao.insert_any(conn, data1)
            dao.insert_any(conn, data2)
    """
    global DB_ENGINE
    return ExTransaction(DB_ENGINE)


class ExConnection(object):
    def __init__(self, engine):
        self.engine = engine

    def __enter__(self):
        self.conn = self.engine.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # TODO 에러처리
        self.conn.close()
        return exc_type is None

    def execute(self, query):
        cur = self.conn.execute(text(query))
        return cur


class ExTransaction(ExConnection):
    def __init__(self, engine):
        super().__init__(engine)
        self.trans = None

    def __enter__(self):
        super().__enter__()
        self.trans = self.conn.begin()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # TODO 에러처리
        if exc_type is None:
            self.trans.commit()
        else:
            self.trans.rollback()
        return super().__exit__(exc_type, exc_val, exc_tb)

