restorepiperun.py

def restore_pipe(p_run: Run, bak_files: Dict[str, str], prev_config: PipelineConfig) -> None:
    """
    Restores the pipeline to the backup files version.

    Args:
        p_run (Run):
            The run model object.
        bak_files (Dict[str, str]):
            Dictionary containing the paths to the .bak files.
        prev_config (PipelineConfig):
            Back up run configuration.

    Returns:
        None
    """
    # check images match
    img_f_list = prev_config["inputs"]["image"]
    if isinstance(img_f_list, dict):
        img_f_list = [
            item for sublist in img_f_list.values() for item in sublist
        ]
    img_f_list = [os.path.basename(i) for i in img_f_list]

    prev_images = pd.read_parquet(
        bak_files['images'], columns=['id', 'name', 'measurements_path']
    )

    if sorted(prev_images['name'].tolist()) != sorted(img_f_list):
        raise CommandError(
            'Images in previous config file does not'
            ' match those found in the previous images.parquet.bak.'
            ' Cannot restore pipeline run.'
        )

    # check forced measurements
    monitor = prev_config["source_monitoring"]["monitor"]
    if monitor:
        forced_parquets = glob(os.path.join(
            p_run.path, 'forced_*.parquet.bak'
        ))

        if not forced_parquets:
            raise CommandError(
                'source_monitoring.monitor is \'True\' in the previous configuration but'
                ' no .bak forced parquet files have been found.'
                ' Cannot restore pipeline run.'
            )
        else:
            # load old associations
            bak_meas_id = pd.read_parquet(
                bak_files['associations'],
                columns=['meas_id']
            )['meas_id'].unique()

            # load backup forced measurements
            forced_meas = pd.concat(
                [pd.read_parquet(i, columns=['id']) for i in forced_parquets]
            )

            # load image meas
            meas = pd.concat(
                [pd.read_parquet(
                    i, columns=['id']
                ) for i in prev_images['measurements_path']]
            )

            # Get forced ids from the associations
            forced_meas_id = bak_meas_id[
                np.isin(bak_meas_id, meas['id'].to_numpy(), invert=True)
            ]

            if not np.array_equal(
                np.sort(forced_meas_id),
                np.sort(forced_meas['id'].to_numpy())
            ):
                raise CommandError(
                    'The forced measurements .bak files do not match the'
                    ' previous run.'
                    ' Cannot restore pipeline run.'
                )

            del meas

    logger.info("Restoring '%s' from backup parquet files.", p_run.name)

    # Delete any new sources
    bak_sources = pd.read_parquet(bak_files['sources'])

    sources_to_delete = (
        Source.objects
        .filter(run=p_run)
        .exclude(id__in=bak_sources.index.to_numpy())
    )

    if sources_to_delete.exists():
        with transaction.atomic():
            n_del, detail_del = sources_to_delete.delete()
            logger.info(
                ('Deleting new sources and associated objects to restore run'
                 ' Total objects deleted: %i'),
                n_del,
            )
            logger.debug('(type, #deleted): %s', detail_del)

    if monitor:
        current_forced_parquets = glob(os.path.join(
            p_run.path, 'forced_*.parquet'
        ))

        current_forced_meas = pd.concat(
            [pd.read_parquet(
                i, columns=['id']
            ) for i in current_forced_parquets]
        )

        ids_to_delete = current_forced_meas.loc[
            ~current_forced_meas['id'].isin(forced_meas['id'].to_numpy()),
            'id'
        ]

        meas_to_delete = Measurement.objects.filter(id__in=ids_to_delete)

        del ids_to_delete

        if meas_to_delete.exists():
            with transaction.atomic():
                n_del, detail_del = meas_to_delete.delete()
                logger.info(
                    ('Deleting forced measurement and associated'
                     ' objects to restore run. Total objects deleted: %i'),
                    n_del,
                )
                logger.debug('(type, #deleted): %s', detail_del)

    # restore source metrics
    logger.info(f'Restoring metrics for {bak_sources.shape[0]} sources.')
    bak_sources = update_sources(bak_sources)

    # remove images from run
    images_to_remove = (
        Image.objects
        .filter(run=p_run)
        .exclude(id__in=prev_images['id'].to_numpy())
    )
    logger.info(f'Removing {len(images_to_remove)} images from the run.')
    if images_to_remove.exists():
        with transaction.atomic():
            p_run.image_set.remove(*images_to_remove)

    # load image meas
    meas = pd.concat(
        [pd.read_parquet(
            i, columns=['id']
        ) for i in prev_images['measurements_path']]
    )

    association_criteria_1 = Q(source_id__in=bak_sources['id'].to_numpy())
    association_criteria_2 = ~Q(meas_id__in=meas['id'].to_numpy())
    associations_to_delete = Association.objects.filter(
        association_criteria_1 and association_criteria_2
    )

    if associations_to_delete.exists():
        with transaction.atomic():
            n_del, detail_del = associations_to_delete.delete()
            logger.info(
                ('Deleting associations to restore run.'
                 ' Total objects deleted: %i'),
                n_del,
            )
            logger.debug('(type, #deleted): %s', detail_del)

    pair_criteria_1 = Q(source_id__in=bak_sources['id'].to_numpy())
    pair_criteria_2 = ~Q(measurement_a__in=meas['id'].to_numpy())
    pair_criteria_3 = ~Q(measurement_b__in=meas['id'].to_numpy())

    pairs_to_delete = MeasurementPair.objects.filter(
       pair_criteria_1 and (pair_criteria_2 | pair_criteria_3)
    )

    if pairs_to_delete.exists():
        with transaction.atomic():
            n_del, detail_del = pairs_to_delete.delete()
            logger.info(
                ('Deleting measurement pairs to restore run.'
                 ' Total objects deleted: %i'),
                n_del,
            )
            logger.debug('(type, #deleted): %s', detail_del)

    logger.info('Restoring run metrics.')
    p_run.n_images = prev_images.shape[0]
    p_run.n_sources = bak_sources.shape[0]
    p_run.n_selavy_measurements = meas.shape[0]
    if monitor:
        p_run.n_forced_measurements = forced_meas.shape[0]

    with transaction.atomic():
        p_run.save()

    # switch files and delete backups
    logger.info('Restoring parquet files and removing .bak files.')
    for i in bak_files:
        bak_file = bak_files[i]
        if i == 'config':
            actual_file = bak_file.replace('.yaml.bak', '_prev.yaml')
        else:
            actual_file = bak_file.replace('.bak', '')
        shutil.copy(bak_file, actual_file)
        os.remove(bak_file)

    if monitor:
        for i in current_forced_parquets:
            os.remove(i)

        for i in forced_parquets:
            new_file = i.replace('.bak', '')
            shutil.copy(i, new_file)
            os.remove(i)

def yesno(question: str) -> bool:
    """
    Simple Yes/No Function.

    Args:
        question (str):
            The question to show to the user for a y/n response.

    Returns:
        True if user enters 'y', False if 'n'.
    """
    prompt = f'{question} ? (y/n): '
    ans = input(prompt).strip().lower()
    if ans not in ['y', 'n']:
        print(f'{ans} is invalid, please try again...')
        return yesno(question)
    if ans == 'y':
        return True
    return False