mirror.py

"""
This module will use evironment variables to download images
from DCAR and push to a container registry. Metadata about
what images to mirror are loaded from `dcar.json`.
"""
import sys
import os
import json
import multiprocessing
import docker

from dcar.request import Session, LoginError
from dcar.image import mirror_image, process_image
from dcar.utils import check_envs, validate_argument

def mirror(dcar_json_input: str = "dcar.json",
           failed_image_output: str = "failed_images.json"):
    """
    Copy from DCAR to container registry.

    Top level function that reads the dcar.json file, downloads
    from DCAR, and pushes to a container registry.

    Environment Variables:
    DCAR_USERNAME (required): Username for authenticating with DCAR.
    DCAR_PASSWORD (required): Password for authenticating with DCAR.
    DCAR_TOTP_SEED (required): Time-based One-time Password seed for authenticating with DCAR.
    REGISTRY_IMAGE_TEMPLATE (required): Container registry name to store images in.
    REGISTRY_USER (required): Username for authenticating with container registry.
    REGISTRY_PASSWORD (required): Password for authenticating with container registry.
    CI_NODE_INDEX (Optional): Index of process mirroring used to slice the list
                              of images being mirrored.
    CI_NODE_TOTAL (Optional): Total number of processes mirroring.
    PUSH_W_DOCKER (Optional): If this environment variable exists, use docker instead of
                              skopeo to push images.
    """

    validate_argument("dcar_json_input", dcar_json_input, str)
    validate_argument("failed_image_output", failed_image_output, str)

    # Make sure envs are set
    check_envs(["DCAR_USERNAME", "DCAR_PASSWORD",
                "DCAR_TOTP_SEED", "REGISTRY_IMAGE_TEMPLATE",
                "REGISTRY_USER", "REGISTRY_PASSWORD"])

    if "CI_NODE_INDEX" in os.environ and "CI_NODE_TOTAL" in os.environ:
        node_index = int(os.environ["CI_NODE_INDEX"])
        node_total = int(os.environ["CI_NODE_TOTAL"])
    else:
        node_index = 1
        node_total = 1

    try:
        with open(dcar_json_input) as json_file:
            dcar_image_dict = json.load(json_file)

        # Since this is probably going to be run as part of a
        # scheduled job, you can define the two environment variables
        # to split the `dcar.json` dictionary for parallel processing
        # This is done to speed up the mirroring images.
        #
        # If you are using GitLab CI, these two environment variables
        # are automatically created if you use the "parallel" tag.
        # - CI_NODE_INDEX
        # - CI_NODE_TOTAL
        complete_dcar_image_list = list(dcar_image_dict.keys())
        dcar_image_list = list(dcar_image_dict.items())[node_index-1::node_total]

        print("processing {}/{} ({}/{} images)".format(node_index, node_total,
                                                        len(dcar_image_list),
                                                        len(complete_dcar_image_list)),
                flush=True)

        # mirror the images
        dcar_failed_image_list = []
        with multiprocessing.Pool(processes=5) as pool:
            #vals = pool.imap(process_image, dcar_image_list)
            vals = pool.imap(process_image, dcar_image_list)
            for image_info in vals:
                dcar_image_name = image_info[0]
                image_exist = image_info[1]
                if not image_exist:
                    dcar_failed_image_list.append(dcar_image_name)

            # write out list of images that failed
            if len(dcar_failed_image_list) > 0:
                # Save the json object with all of the failed images.
                with open(failed_image_output, 'w') as outfile:
                    json.dump(dcar_failed_image_list, outfile)
                if len(dcar_failed_image_list) != 0:
                    print("ERROR: The following images failed to mirror", flush=True, file=sys.stderr)
                    for failed_image in dcar_failed_image_list:
                        print("- " + failed_image, flush=True, file=sys.stderr)
    except LoginError as exception:
        sys.exit("FATAL: {}".format(exception))