taskcluster/taskgraph/util/taskcluster.py
author Wes Kocher <wkocher@mozilla.com>
Mon, 11 Sep 2017 17:35:09 -0700
changeset 380171 bda524beac249b64aa36016800502a34073bf35a
parent 379989 01bf0c29331e6bbee4a7bf674010ea1ebc41cb55
child 381495 a43387ee65d7f889aac9b2581f12a6f2512e59b6
permissions -rw-r--r--
Backed out 2 changesets (bug 1342392) for bustage a=backout Backed out changeset b0e86286b173 (bug 1342392) Backed out changeset 01bf0c29331e (bug 1342392) MozReview-Commit-ID: 71fcOZ2RgEp

# -*- coding: utf-8 -*-

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from __future__ import absolute_import, print_function, unicode_literals

import functools
import yaml
import requests
from mozbuild.util import memoize
from requests.packages.urllib3.util.retry import Retry
from requests.adapters import HTTPAdapter


@memoize
def get_session():
    session = requests.Session()
    retry = Retry(total=5, backoff_factor=0.1,
                  status_forcelist=[500, 502, 503, 504])
    session.mount('http://', HTTPAdapter(max_retries=retry))
    session.mount('https://', HTTPAdapter(max_retries=retry))
    return session


def _do_request(url):
    session = get_session()
    response = session.get(url, stream=True)
    if response.status_code >= 400:
        # Consume content before raise_for_status, so that the connection can be
        # reused.
        response.content
    response.raise_for_status()
    return response


def _handle_artifact(path, response):
    if path.endswith('.json'):
        return response.json()
    if path.endswith('.yml'):
        return yaml.load(response.text)
    response.raw.read = functools.partial(response.raw.read,
                                          decode_content=True)
    return response.raw


def get_artifact_url(task_id, path, use_proxy=False):
    if use_proxy:
        ARTIFACT_URL = 'http://taskcluster/queue/v1/task/{}/artifacts/{}'
    else:
        ARTIFACT_URL = 'https://queue.taskcluster.net/v1/task/{}/artifacts/{}'
    return ARTIFACT_URL.format(task_id, path)


def get_artifact(task_id, path, use_proxy=False):
    """
    Returns the artifact with the given path for the given task id.

    If the path ends with ".json" or ".yml", the content is deserialized as,
    respectively, json or yaml, and the corresponding python data (usually
    dict) is returned.
    For other types of content, a file-like object is returned.
    """
    response = _do_request(get_artifact_url(task_id, path, use_proxy))
    return _handle_artifact(path, response)


def list_artifacts(task_id, use_proxy=False):
    response = _do_request(get_artifact_url(task_id, '', use_proxy).rstrip('/'))
    return response.json()['artifacts']


def get_index_url(index_path, use_proxy=False):
    if use_proxy:
        INDEX_URL = 'http://taskcluster/index/v1/task/{}'
    else:
        INDEX_URL = 'https://index.taskcluster.net/v1/task/{}'
    return INDEX_URL.format(index_path)


def find_task_id(index_path, use_proxy=False):
    response = _do_request(get_index_url(index_path, use_proxy))
    return response.json()['taskId']


def get_artifact_from_index(index_path, artifact_path, use_proxy=False):
    full_path = index_path + '/artifacts/' + artifact_path
    response = _do_request(get_index_url(full_path, use_proxy))
    return _handle_artifact(full_path, response)


def get_task_url(task_id, use_proxy=False):
    if use_proxy:
        TASK_URL = 'http://taskcluster/queue/v1/task/{}'
    else:
        TASK_URL = 'https://queue.taskcluster.net/v1/task/{}'
    return TASK_URL.format(task_id)


def get_task_definition(task_id, use_proxy=False):
    response = _do_request(get_task_url(task_id, use_proxy))
    return response.json()