path: root/library/recipes/docker_compose/lib/__init__.py

import os
import random
import re
import subprocess
import six

import yaml
import logging
import argparse
import yatest.common

from test import const
from six.moves import input
import library.python.fs as fs
import library.python.testing.recipe


logger = logging.getLogger(__name__)


class DockerComposeRecipeException(Exception):

    def __init__(self, msg):
        super(DockerComposeRecipeException, self).__init__("[[bad]]{}[[rst]]".format(msg))


def start(argv):
    args = _parse_args(argv)

    yml_file, cwd = get_compose_file_and_cwd(argv)
    _verify_compose_file(yml_file)

    env = _setup_env()

    docker_compose = get_docker_compose()

    recipe_config = _get_recipe_config(args)

    deprecated_context = _get_docker_deprecated_context(args)

    docker_context = recipe_config.get('context', deprecated_context)
    if docker_context:
        context_map = _create_context(docker_context, os.path.dirname(yml_file), yatest.common.work_path("docker_context_root"))
        env.update(context_map)

    test_host_name = _get_test_host(args, recipe_config)
    if test_host_name:
        yml_file = _setup_test_host(test_host_name, yml_file, env)
        library.python.testing.recipe.set_env("DONT_CREATE_TEST_PROCESS_GROUP", "1")  # XXX find out why docker-compose hangs in other case
        library.python.testing.recipe.set_env(
            "TEST_COMMAND_WRAPPER", " ".join([docker_compose, "-f", yml_file, "exec", "-T", test_host_name])
        )

    library.python.testing.recipe.set_env("DOCKER_COMPOSE_FILE", yml_file)

    networks = _get_networks(recipe_config)
    if networks:
        subnets = set()

        for net_name, net_config in six.iteritems(networks):
            while True:
                subnet = "fc00:420:%04x::/32" % random.randrange(16**4)
                if subnet not in subnets:
                    break

            subnets.add(subnet)

            net_args = ["docker", "network", "create", "--subnet", subnet]
            if net_config and net_config.get("ipv6", False):
                net_args.append("--ipv6")
            net_args.append(net_name)

            yatest.common.execute(net_args, cwd=cwd)

    yatest.common.execute([docker_compose, "-f", yml_file, "--log-level", "DEBUG", "--no-ansi", "up", "-d", "--build",
                           "--force-recreate"],
                          cwd=cwd, env=env)


def stop(argv):
    if yatest.common.get_param("docker-pause"):
        library.python.testing.recipe.tty()
        try:
            input("\ndocker_compose will stop, press <Enter> to continue")
        except KeyboardInterrupt:
            pass

    docker_compose = get_docker_compose()
    yaml, cwd = get_compose_file_and_cwd(argv)
    args = _parse_args(argv)
    recipe_config = _get_recipe_config(args)
    failed_containers = []
    containers = _get_containers(yaml)

    try:
        containers_ids_res = yatest.common.execute([docker_compose, "-f", yaml, "--log-level", "DEBUG", "--no-ansi", "ps", "-q"], cwd=cwd, stdout=subprocess.PIPE, text=True)
        if containers_ids_res.exit_code != 0:
            raise DockerComposeRecipeException("'docker-compose ps' returned {}'".format(containers_ids_res.exit_code))

        containers_ids = str.splitlines(containers_ids_res.std_out)
        if len(containers_ids) == 0:
            raise DockerComposeRecipeException("'docker-compose ps' output is empty '{}'".format(containers_ids_res.std_out))

        for container_id in containers_ids:
            container_id_status_res = yatest.common.execute(["docker", "ps", "-a", "--filter", "id=" + container_id, "--format", "{{.Status}}\t{{.Names}}"], cwd=cwd)

            if container_id_status_res.exit_code != 0:
                raise DockerComposeRecipeException("'docker ps' returned {}'".format(container_id_status_res.exit_code))

            status_line, container_name = six.ensure_str(container_id_status_res.std_out).split('\t')
            if "Up" in status_line or "Exited (0)" in status_line:
                continue
            failed_containers.append(container_name)

        yatest.common.execute([docker_compose, "-f", yaml, "--log-level", "DEBUG", "--no-ansi", "stop"], cwd=cwd)

        _dump_container_logs(containers, _get_requested_paths(yaml, recipe_config))

    finally:
        if not yatest.common.context.test_debug:
            yatest.common.execute([get_docker_compose(), "-f", yaml, "--log-level", "DEBUG", "--no-ansi", "down", "--rmi", "local"], cwd=cwd)

        networks = _get_networks(recipe_config)
        if networks:
            for net_name in six.iterkeys(networks):
                yatest.common.execute(["docker", "network", "rm", net_name], cwd=cwd)

        if len(failed_containers) > 0:
            raise DockerComposeRecipeException("Has failed containers: {}".format(", ".join(failed_containers)))


def get_compose_file_and_cwd(args):
    if "DOCKER_COMPOSE_FILE" in os.environ:
        yaml_file = os.environ["DOCKER_COMPOSE_FILE"]
    else:
        args = _parse_args(args)
        if args.compose_file and args.compose_file != "$DOCKER_COMPOSE_FILE":
            yaml_file = yatest.common.source_path(args.compose_file)
        else:
            yaml_file = yatest.common.test_source_path("docker-compose.yml")
    return yaml_file, os.path.dirname(yaml_file)


def get_docker_compose():
    docker_compose = yatest.common.build_path("library/recipes/docker_compose/bin/docker-compose")
    if not os.path.exists(docker_compose):
        raise DockerComposeRecipeException("cannot find docker_compose by build_path '{}'".format(docker_compose))
    os.chmod(docker_compose, 0o755)
    return docker_compose


def _parse_args(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument("--recipe-config-file", help="Path recipe config yml file (Arcadia related)")
    parser.add_argument("--compose-file", help="Path to docker-compose.yml file (Arcadia related)")
    parser.add_argument("--context-file", help="Path to docker-context.yml file (Arcadia related)")
    parser.add_argument("--test-host", help="Name of service in docker-compose file that will host test execution")
    return parser.parse_args(argv)


def _get_recipe_config(args):
    if args.recipe_config_file and args.recipe_config_file != "$RECIPE_CONFIG_FILE":
        config_path = yatest.common.source_path(args.recipe_config_file)
        if not os.path.exists(config_path):
            raise DockerComposeRecipeException("Cannot find specified recipe config file '{}'".format(args.recipe_config_file))
        with open(config_path) as f:
            return yaml.load(f, Loader=yaml.FullLoader)
    return {}


def _get_docker_deprecated_context(args):
    # XXX: to be removed when all recipes use new config file
    if args.context_file and args.context_file != "$DOCKER_CONTEXT_FILE":
        context_file = yatest.common.source_path(args.context_file)
        if not os.path.exists(context_file):
            raise DockerComposeRecipeException("Cannot find context file by {}".format(context_file))
        with open(context_file) as f:
            return yaml.load(f, Loader=yaml.FullLoader)
    return None


def _setup_env():
    env = os.environ.copy()
    env["CURRENT_USER"] = "{}:{}".format(os.getuid(),  os.getgid())
    # Setup extra env.vars. to be able to pass coverage dir to the docker
    for name in const.COVERAGE_ENV_VARS:
        if name in env:
            env["{}_DIRNAME".format(name.rsplit('_', 1)[0])] = os.path.dirname(env[name])
    return env


def _create_context(context, init_dir, context_root):

    def copy_files(src, dst):
        if os.path.isfile(src):
            fs.copy_file(src, dst)
        else:
            if not os.path.exists(dst):
                os.makedirs(dst)

            for name in os.listdir(src):
                copy_files(os.path.join(src, name), os.path.join(dst, name))

    context_map = {}

    fs.ensure_dir(context_root)
    for context_name in context:
        _verify_context_name(context_name)
        context_dir = os.path.join(context_root, context_name)
        copy_files(init_dir, context_dir)
        for item in context[context_name]:
            if len(item.keys()) != 1 or len(item.values()) != 1:
                raise DockerComposeRecipeException("Context item should be in form of <source>:<destination> item")
            source_path, target_path = next(iter(six.iteritems(item)))
            if source_path.startswith("build://"):
                source_path = yatest.common.build_path(source_path[len("build://"):])
            elif source_path.startswith("arcadia://"):
                source_path = yatest.common.source_path(source_path[len("arcadia://"):])
            else:
                raise DockerComposeRecipeException("Source path should start with 'build://' or 'arcadia://'")
            target_path = os.path.join(context_dir, target_path.lstrip("/"))
            fs.ensure_dir(os.path.dirname(target_path))
            copy_files(source_path, target_path)

        context_map[context_name] = context_dir

    return context_map


def _verify_context_name(name):
    assert re.match("^[a-zA-Z0-9]+$", name), "Context name '{}' has incorrect symbols".format(name)


def _verify_compose_file(file_path):
    with open(file_path) as f:
        data = yaml.safe_load(f)
    known_images = set()
    for service, settings in six.iteritems(data["services"]):
        if "image" in settings:
            image = settings["image"]
            if "build" not in settings and "@sha256" not in image and image not in known_images:
                message = "Using image without specified sha256 (e.g. redis:alpine@sha256:66ccc75f079ab9059c900e9545bbd271bff78a66f94b45827e6901f57fb973f1) and build section is not supported"
                logger.error(message)
                raise DockerComposeRecipeException(message)
            known_images.add(image)


def _get_networks(config):
    return config.get("networks")


def _get_test_host(args, config):
    def get_from_args():
        if args.test_host and args.test_host != "$DOCKER_TEST_HOST":
            return args.test_host
        return None
    return config.get('test-host', get_from_args())


def _setup_test_host(test_host_name, yml_path, env):
    with open(yml_path) as f:
        data = yaml.load(f, Loader=yaml.FullLoader)

    for service, settings in six.iteritems(data["services"]):
        if service == test_host_name:
            overwritten_yml_path = yatest.common.work_path("docker_compose_for_test.yml")
            if "env" not in settings:
                settings["environment"] = []
            for env_key, env_value in six.iteritems(env):
                settings["environment"].append("{}={}".format(env_key, env_value))

            if "volumes" not in settings:
                settings["volumes"] = []

            if "command" in settings:
                raise DockerComposeRecipeException("Test hosting service '{}' has `command` section which is not supported by testing framework".format(test_host_name))
            settings["command"] = "sleep 3600"

            settings["tty"] = False

            if "user" in settings:
                raise DockerComposeRecipeException("Test hosting service '{}' has `user` section which is not supported by testing framework".format(test_host_name))
            settings["user"] = "$CURRENT_USER"

            if "build" in settings:
                conext = settings["build"].get("context")
                if conext == ".":
                    settings["build"]["context"] = os.path.dirname(yml_path)

            bind_paths = [
                yatest.common.build_path(),
                os.environ.get("ORIGINAL_SOURCE_ROOT", yatest.common.source_path()),
            ]

            if yatest.common.runtime.context.test_tool_path:
                bind_paths.append(yatest.common.runtime.context.test_tool_path)

            for k in [
                "PORT_SYNC_PATH",
                "OS_SDK_ROOT_RESOURCE_GLOBAL",
                "LLD_ROOT_RESOURCE_GLOBAL",
                "ASAN_SYMBOLIZER_PATH",
                "LSAN_SYMBOLIZER_PATH",
                "MSAN_SYMBOLIZER_PATH",
                "UBSAN_SYMBOLIZER_PATH",
                "TMPDIR",
            ]:
                if k in os.environ:
                    p = os.environ[k]
                    if p not in bind_paths and os.path.exists(p):
                        bind_paths.append(p)

            for bind_path in bind_paths:
                settings["volumes"].append({
                    "type": "bind",
                    "source": bind_path,
                    "target": bind_path,
                })
                real_bind_path = os.path.realpath(bind_path)
                if real_bind_path != bind_path:
                    settings["volumes"].append({
                        "type": "bind",
                        "source": real_bind_path,
                        "target": real_bind_path,
                    })

            with open(overwritten_yml_path, "w") as f:
                yaml.dump(data, f)

            return overwritten_yml_path

    raise DockerComposeRecipeException("Service with name '{}' was not found to be setup as a host for running test".format(test_host_name))


def _dump_container_logs(containers, requested_paths):
    # add links to container logs
    container_logs = yatest.common.output_path("containers")
    if not os.path.exists(container_logs):
        os.makedirs(container_logs)

    for container_id, container_name in six.iteritems(containers):
        try:
            output_path = os.path.join(container_logs, container_name)
            os.makedirs(output_path)
            for output_type in ['std_out', 'std_err']:
                output_log_path = os.path.join(output_path, "container_{}.log".format(output_type))
                res = yatest.common.execute(["docker", "logs", container_id])
                with open(output_log_path, "w") as f:
                    f.write(six.ensure_str(getattr(res, output_type)))

            for path in requested_paths.get(container_id, []):
                try:
                    yatest.common.execute(["docker", "cp", "-L", "{}:{}".format(container_id, path), output_path])
                except yatest.common.ExecutionError:
                    logging.exception("Error while copying %s from %s", path, container_name)

        except Exception:
            logger.exception("Error collecting container's log with name {} and id {}".format(container_name, container_id))


def _get_requested_paths(yaml_path, config):
    requested_logs = {}
    if config:
        for service_name, logs in six.iteritems(config.get("save", {})):
            try:
                res = yatest.common.execute([get_docker_compose(), "-f", yaml_path, "ps", "-q", service_name])
                requested_logs[six.ensure_str(res.std_out).strip()] = logs
            except yatest.common.ExecutionError:
                logging.exception("Error while trying to find docker compose service by name: %s", service_name)
    return requested_logs


def _get_containers(yaml_path):
    res = yatest.common.execute([get_docker_compose(), "-f", yaml_path, "ps", "-q"])
    return {container_id: _get_container_name(container_id) for container_id in filter(None, six.ensure_str(res.std_out).split("\n"))}


def _get_container_name(container_id):
    res = yatest.common.execute(["docker", "inspect", "--format={{.Name}}", container_id])
    container_name = six.ensure_str(res.std_out).strip("/").strip()
    return container_name