yuxin_pms/utils.py

import yaml
from sqlalchemy import text


def load_config(env='test'):
    """ Loads the configuration file
    """
    with open('config.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    return config[env]


def truncate_target_db(conn, target_db) -> None:
    """ Truncates the target table
    """
    query = """
        SELECT count(1) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '{target_db}'
        """.format(target_db=target_db)
    if conn.execute(text(query)).scalar() == 1:
        sql = """
                TRUNCATE TABLE {target_db}
            """.format(target_db=target_db)
        conn.execute(text(sql))