download_sys_perf_binaries.py

#!/usr/bin/env python3
"""Download the binaries from a previous sys-perf run."""

import argparse

import requests

BASE_URI = "https://evergreen.mongodb.com/rest/v2/"


def _get_auth_headers(evergreen_api_user, evergreen_api_key):
    return {
        "Api-User": evergreen_api_user,
        "Api-Key": evergreen_api_key,
    }


def _get_build_id(build_variant_name, version_id, auth_headers):
    url = BASE_URI + "versions/" + version_id
    response = requests.get(url, headers=auth_headers)
    if response.status_code != 200:
        raise ValueError("Invalid version_id:", version_id)

    version_json = response.json()
    build_variants = version_json["build_variants_status"]
    for build_variant in build_variants:
        if build_variant["build_variant"] == build_variant_name:
            return build_variant["build_id"]

    raise RuntimeError(
        "The compile-variant "
        + build_variant_name
        + " does not exist for the build with version_id "
        + version_id
    )


# All of our sys-perf compile variants have exactly one task,
# so we can safely select the first one from the build variants tasks.
def _get_task_id(build_id, auth_headers):
    url = BASE_URI + "builds/" + build_id
    response = requests.get(url, headers=auth_headers)
    if response.status_code != 200:
        raise RuntimeError("Unexpected error when trying to reach" + url)

    build_json = response.json()
    task_list = build_json["tasks"]
    if len(task_list) != 1:
        raise RuntimeError("Recieved unexpected tasklist:", task_list)
    return task_list[0]


# The API used here always grabs the latest execution
def _get_binary_details(task_id, auth_headers):
    url = BASE_URI + "tasks/" + task_id
    response = requests.get(url, headers=auth_headers)
    if response.status_code != 200:
        raise RuntimeError("Unexpected error when trying to reach" + url)

    task_json = response.json()
    if task_json["status"] != "success":
        raise RuntimeError("The task " + task_id + " did not run sucessfully")

    # The binary will always be the first artifact, unless we make large changes to system_perf.yml
    artifacts = task_json["artifacts"]
    if len(artifacts) > 0 and artifacts[0]["name"].startswith("mongo"):
        return artifacts[0]
    raise RuntimeError("Unexpected list of artifacts:" + artifacts)


def _get_binary_url(version_id, build_variant, evergreen_api_user, evergreen_api_key):
    auth_headers = _get_auth_headers(evergreen_api_user, evergreen_api_key)
    build_id = _get_build_id(build_variant, version_id, auth_headers)
    task_id = _get_task_id(build_id, auth_headers)
    binary_json = _get_binary_details(task_id, auth_headers)
    return binary_json["url"]


def _download_binary_file(url, save_path):
    response = requests.get(url, stream=True)
    if response.status_code == 200:
        with open(save_path, "wb") as file:
            for chunk in response.iter_content(chunk_size=1024):
                file.write(chunk)
    else:
        raise RuntimeError("Failed to download the file " + url)


def _download_sys_perf_binaries(version_id, build_variant, evergreen_api_user, evergreen_api_key):
    url = _get_binary_url(version_id, build_variant, evergreen_api_user, evergreen_api_key)
    _download_binary_file(url, "binary.tar.gz")


if __name__ == "__main__":
    argParser = argparse.ArgumentParser()
    argParser.add_argument(
        "-v", "--version_id", help="Evergreen version_id from which binaries will be downloaded"
    )
    argParser.add_argument(
        "-b", "--build_variant", help="Build variant for which binaries will be downloaded"
    )
    argParser.add_argument(
        "-u",
        "--evergreen_api_user",
        help="Evergreen API user, see https://spruce.mongodb.com/preferences/cli",
    )
    argParser.add_argument(
        "-k",
        "--evergreen_api_key",
        help="Evergreen API key, see https://spruce.mongodb.com/preferences/cli",
    )
    args = argParser.parse_args()

    _download_sys_perf_binaries(
        args.version_id, args.build_variant, args.evergreen_api_user, args.evergreen_api_key
    )