"""
MySQL Connector.

Manages interactions with MySQL
for ows_transcoding database.
"""

from contextlib import contextmanager
from typing import Generator

from sqlalchemy import create_engine, pool, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, declarative_base, sessionmaker

from transcoding import config


def _create_engine(db_url: str) -> Engine:
    """Create engine based on config settings."""
    if config.POOL_CLASS == pool.QueuePool:
        return create_engine(
            db_url,
            pool_size=config.POOL_SIZE,
            max_overflow=config.POOL_MAX_OVERFLOW,
            pool_recycle=config.POOL_RECYCLE_MS,
        )
    return create_engine(db_url, poolclass=config.POOL_CLASS)


# please don't use the following private variables directly;
# use db_session
_db_engine = _create_engine(config.DB_URL)
_db_session = sessionmaker(bind=_db_engine)

BaseModel = declarative_base()


@contextmanager
def db_session(
    turn_off_foreign_key_checks: bool = False,
) -> Generator[Session, None, None]:
    """Provide a transactional scope around a series of operations.

    Taken from http://docs.sqlalchemy.org/en/latest/orm/session_basics.html.
    This handles rollback and closing of session, so there is no need
    to do that throughout the code.
    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = _db_session()
    try:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
        session.close()