import datetime

from dateutil.relativedelta import relativedelta

from service import trans_service
from utils.log.trans_log import logger

if __name__ == '__main__':

    datas = trans_service.get_all_partitioned_tables()

    now_month = datetime.datetime.now()
    next_month = now_month + relativedelta(months=1)
    pname = f'p{str(now_month.year) + str(now_month.month).zfill(2)}'
    date_str = f'{str(next_month.year)}-{str(next_month.month).zfill(2)}-01'
    for data in datas:
        trans_service.add_partition(data['TABLE_NAME'], pname, date_str)

    logger.info("添加分区成功")

    save_month_dict = trans_service.get_sys_conf('online_data_save_month')

    table_types = set()

    for data in datas:
        table_name = data['TABLE_NAME']
        wind_factory = table_name.split('_')[0]
        table_type = table_name.split('_')[1]
        table_types.add(table_type)
        del_month = save_month_dict.get(table_type, 12)
        pmonth = datetime.datetime.now() - relativedelta(months=del_month + 1)
        trans_service.delelet_partition(table_name, pmonth)

    for table_type in table_types:
        del_month = save_month_dict.get(table_type, 12)
        pmonth = datetime.datetime.now() - relativedelta(months=del_month)
        exists_date = f'{pmonth.year}-{pmonth.month}-01'
        trans_service.update_expired_data(table_type,exists_date)

    logger.info("删除过期分区成功")
