assets/__init__.py

# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
# or more contributor license agreements. Licensed under the Elastic License;
# you may not use this file except in compliance with the Elastic License.

import os
from pathlib import Path

assets_dir = Path(__file__).parent
branches = ("production", "staging", "snapshot")


def walk():
    """
    Traverse all the local assets.

    :return: generator yielding (branch, package, version) of all the local assets
    """

    for branch in branches:
        for asset_branch, packages, _ in os.walk(assets_dir / branch):
            for package in packages:
                if package.startswith("."):
                    continue
                for _, versions, _ in os.walk(Path(asset_branch) / package):
                    for version in versions:
                        yield branch, package, version
                    break
            break


def get_meta(branch, package, version):
    """
    Get the meta-data of a local asset.

    :param branch: one among 'production', 'staging', 'snapshot'
    :param package: package name, ex. 'endpoint'
    :param version: package version, ex. '8.3.0'
    :return: dictionary containing the meta-data
    """

    import yaml

    meta_filename = assets_dir / branch / package / version / "meta.yml"
    if meta_filename.exists():
        with open(meta_filename) as f:
            return yaml.safe_load(f)


def get_local_assets(package, path):
    """
    Retrieve the list of a package's local assets.

    :param package: name and version of the package, ex. 'endpoint/8.3.0'
    :param path: path on disk searched for the assets
    :return: generator yielding (path, content) pairs as they are traversed
    """

    saved_cwd = os.getcwd()
    os.chdir(path)
    try:
        if not Path(package).exists():
            raise ValueError(f"Package not found: {package}")
        for root, _, files in os.walk(package):
            for file in files:
                with open(Path(root) / file, "rb") as f:
                    yield f.name, f.read()
    finally:
        os.chdir(saved_cwd)


def get_remote_assets(package, repo):
    """
    Retrieve the list of a package's remote assets.

    :param package: name and version of the package, ex. 'endpoint/8.3.0'
    :param repo: repository object searched for the assets
    :return: generator yielding the remote assets entries
    """

    from github import GithubException

    for branch in branches:
        try:
            entries = repo.get_contents(package, ref=branch)
        except GithubException:
            continue

        while entries:
            entry = entries.pop(0)
            if entry.type == "dir":
                entries += repo.get_contents(entry.path, ref=branch)
            else:
                yield entry
        return

    raise ValueError(f"Package not found: {package}")


def download_assets(entries):
    """
    Download the assets of a package.

    :param entries: assets entries as generated by :py:func:`.get_assets`
    :return: generator yielding (path, content) pairs as they get ready
    """

    from requests_futures.sessions import FuturesSession
    from concurrent.futures import as_completed

    session = FuturesSession()
    futures = []

    for entry in entries:
        future = session.get(entry.download_url)
        future.entry = entry
        futures.append(future)

    for future in as_completed(futures):
        res = future.result()
        res.raise_for_status()
        yield future.entry.path, res.content

    session.close()