yuxin_pms/utils.py

import yaml
from sqlalchemy import text
from sqlalchemy.types import Date, INTEGER


def load_config(env='test'):
    """ Loads the configuration file
    """
    with open('config.yaml', 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    
    print('loading config for env: {}'.format(env))
    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))


def find_renewal_contract(conn, contract_id):
    query = """
        select crc1.id, crc1.begin_time, crc1.end_time, crc1.quite_date, 
               crc2.id, crc2.begin_time, crc2.end_time, crc2.quite_date, 
               crc3.id, crc3.begin_time, crc3.end_time, crc3.quite_date
        from yuxin_contract.cont_renter_contract crc1
        left join yuxin_contract.cont_renter_contract crc2 on crc2.contract_pid=crc1.id and crc2.is_delete=0 and crc2.contract_status<>4 and crc2.sign_type=2
        left join yuxin_contract.cont_renter_contract crc3 on crc3.contract_pid=crc2.id and crc3.is_delete=0 and crc3.contract_status<>4 and crc3.sign_type=2
        where crc1.is_delete=0 and crc1.contract_status<>4 and crc1.sign_type=2
            and crc1.contract_pid='{contract_id}'
        """.format(contract_id=contract_id)

    result = conn.execute(text(query)).fetchone()
    if result is None:
        return None

    if result[11] is not None:
        return result[11]
    elif result[10] is not None:
        return result[10]
    elif result[7] is not None:
        return result[7]
    elif result[6] is not None:
        return result[6]
    elif result[3] is not None:
        return result[3]
    elif result[2] is not None:
        return result[2]
    return None


def update_column_comment(conn, target_db, cols):
    for key, value in cols.items():
        dtype = value[0]
        if dtype == Date:
            dtype = 'date'
        elif dtype == INTEGER:
            dtype = 'INT'
        conn.execute(text('alter table {target_db} modify column `{col}` {dtype} comment \'{comment}\''
                          .format(target_db=target_db, col=key, dtype=dtype, comment=value[1])))