"""db_operations.py.

Database level utility class for testing against the mkt_program_info schema.
"""

import sys

from marketing import config
from marketing.connectors import mysql
from marketing.models import highlight


DROP_MKT_PROGRAM_INFO = """
DROP TABLE IF EXISTS `mkt_program_info`;
"""


@mysql.autosession()
def create_mkt_program_info_table(session):
    """Create the mkt_program_info table.

    Args:
        session (Session): the mysql session.
    """
    _exit_if_not_test_environment(session)
    highlight.Highlight.__table__.create(mysql._db_engine)


@mysql.autosession()
def drop_mkt_program_info_table(session):
    """Drop the mkt_program_info table.

    Args:
        session (Session): the mysql session.
    """
    _exit_if_not_test_environment(session)
    session.execute(DROP_MKT_PROGRAM_INFO)


@mysql.autosession()
def seed_models(models, session):
    """Save the given model(s) to the DB.

    Args:
        models (mysql.BaseModel): the base model.
        session (Sesssion): the mysql session.
    """
    if not hasattr(models, '__iter__'):
        models = [models]

    _exit_if_not_test_environment(session)
    for model in models:
        session.merge(model)


def _exit_if_not_test_environment(session):
    """For safety, only run tests in test environment pointed to sqlite.

    Exit immediately if not in test environment or not pointed to sqlite.
    """
    if config.ENVIRONMENT != config.TEST_ENVIRONMENT:
        sys.exit('Environment must be set to {}.'.format(
            config.TEST_ENVIRONMENT))
    if 'sqlite' not in session.bind.url.drivername:
        sys.exit('Tests must point to sqlite database.')