update_existing_rows_drsstep.py 文件源码

def main(config):
    '''
    After we added the new columns drsstep and roi,
    existing rois contained nulls. This script updates those rows
    with the correct values.
    '''

    config = load_config(config)

    database.init(**config['processing_database'])
    database.connect()

    engine = create_mysql_engine(**config['fact_database'])
    with engine.connect() as conn:
        df = pd.read_sql_table('RunInfo', conn, columns=[
            'fNight', 'fRunID', 'fDrsStep', 'fROI'
        ])

    df.set_index(['fNight', 'fRunID'], inplace=True)

    query = RawDataFile.select().where(RawDataFile.roi == None)
    for raw_data_file in tqdm(query, total=query.count()):

        night = raw_data_file.night.year * 10000 + raw_data_file.night.month * 100 + raw_data_file.night.day

        raw_data_file.roi = df.loc[(night, raw_data_file.run_id), 'fROI']
        raw_data_file.save()

    query = DrsFile.select().where(DrsFile.roi == None)
    for drs_file in tqdm(query, total=query.count()):

        night = drs_file.night.year * 10000 + drs_file.night.month * 100 + drs_file.night.day

        drs_file.roi = df.loc[(night, drs_file.run_id), 'fROI']
        drs_file.drs_step = df.loc[(night, drs_file.run_id), 'fDrsStep']
        drs_file.save()