list-artifact-sizes.py

#!/usr/bin/env python

from __future__ import print_function

import click
import json
import logging
import requests

from collections import defaultdict, namedtuple
from concurrent.futures import ThreadPoolExecutor

INDEX = 'https://index.taskcluster.net/v1/'
QUEUE = 'https://queue.taskcluster.net/v1/'

logger = logging.getLogger('main')


IndexTask = namedtuple('IndexTask', ['task_id', 'namespace'])


def info(msg, *args, **kwargs):
    logger.info(msg.format(*args, **kwargs))


def index_namespaces(namespace, limit=1000):
    r = requests.post(INDEX + 'namespaces/' + namespace,
                      json={'limit': limit})
    for n in r.json()['namespaces']:
        yield n['namespace']


def index_tasks(namespace, limit=1000):
    r = requests.post(INDEX + 'tasks/' + namespace,
                      json={'limit': limit})
    for t in r.json()['tasks']:
        yield IndexTask(t['taskId'], t['namespace'])


def tasks_by_changeset(revisions_limit):
    for n in index_namespaces('gecko.v2.mozilla-central.nightly.revision', revisions_limit):
        for t in index_tasks(n + '.firefox'):
            yield t


def list_artifacts(task_id):
    r = requests.get(QUEUE + 'task/%s/artifacts' % task_id)
    if r.status_code != 200:
        return []
    return [a['name'] for a in r.json()['artifacts']]


def gecko_branches():
    for n in index_namespaces('gecko.v2'):
        yield n


def products(branch):
    for n in index_namespaces('gecko.v2.{}.latest'.format(branch)):
        yield n


def all_build_artifacts_iter(ns):
    for task in index_tasks(ns):
        info('build type: {}, task: {}', task.namespace, task.task_id)
        for a in list_artifacts(task.task_id):
            yield (task, a)


def get_artifact_size(task, artifact):
    url = '{}task/{}/artifacts/{}'.format(QUEUE, task.task_id, artifact)
    info('artifact: {}', url)
    r = requests.head(url, allow_redirects=True)
    if r.status_code != 200:
        return 0
    return int(r.headers.get('Content-length', '0'))


def all_build_artifacts(ns):
    tasks = defaultdict(lambda: {'id': '', 'artifacts': []})
    with ThreadPoolExecutor(max_workers=16) as executor:
        for task, artifact, size in executor.map(lambda i: (i[0], i[1], get_artifact_size(*i)),
                                                 all_build_artifacts_iter(ns)):
            tasks[task.namespace]['id'] = task.task_id
            tasks[task.namespace]['artifacts'].append((artifact, size))
    return tasks



@click.command()
@click.argument('branch')
@click.argument('output')
@click.option('--product', default='firefox', help='Product for which to list artifacts')
def list_build_artifacts(branch, product, output):
    ns = 'gecko.v2.{}.latest.{}'.format(branch, product)
    all_artifacts = all_build_artifacts(ns)
    with open(output, 'wb') as f:
        json.dump(all_artifacts, f)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    list_build_artifacts()