recipe_modules/luci_config/api.py

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import base64

from google.protobuf import text_format as textpb

from recipe_engine import recipe_api

from PB.go.chromium.org.luci.buildbucket.proto import project_config as bb_pb2
from PB.go.chromium.org.luci.cv.api.config.v2 import config as cv_config_pb2
from PB.go.chromium.org.luci.milo.api.config import project as milo_pb2


class LuciConfigApi(recipe_api.RecipeApi):
    """Module for polling and parsing luci config files via the luci-config API."""

    def fetch_config(self, config_name, message_type, project=None, local_dir=None):
        """Fetch a config file from the luci-config API and return it as a proto object.

        Args:
            config_name (str): The name of the config file to fetch, e.g.
                "commit-queue.cfg".
            message_type (type): The Python type corresponding to the
                config's protobuf message type.
            project (str): The name of the LUCI project to fetch the config
                from; e.g. "fuchsia". Defaults to the project that the
                current Buildbucket build is running in.
            local_dir (Path): If specified, assumed to point to a local
                directory of files generated by lucicfg. The specified config
                file will be read from the corresponding local file rather
                than fetching it from the LUCI Config service.
        """
        if not project:
            project = self.m.buildbucket.build.builder.project
            assert project, "buildbucket input has no project set"

        if local_dir:
            text = self.m.file.read_text(
                "read %s %s" % (project, config_name), local_dir.join(config_name)
            )
        else:
            with self.m.step.nest("fetch %s %s" % (project, config_name)):
                text = self._fetch_config_textproto(project, config_name)

        cfg = message_type()
        textpb.Parse(text, cfg)
        return cfg

    def _fetch_config_textproto(self, project, config_name):
        url = (
            "https://luci-config.appspot.com/_ah/api/config/v1/config_sets/projects/%s/config/%s"
            % (project, config_name)
        )
        token = self.m.service_account.default().get_access_token()
        response = self.m.url.get_json(
            url,
            step_name="get",
            headers={"Authorization": "Bearer %s" % token},
            default_test_data={"content": ""},
        )
        return base64.b64decode(response.output["content"])

    def buildbucket(self, **kwargs):
        return self.fetch_config("cr-buildbucket.cfg", bb_pb2.BuildbucketCfg, **kwargs)

    def commit_queue(self, config_name=None, **kwargs):
        # Support loading a CQ config file with a non-default name to support
        # projects that don't want a dedicated CQ instance but for which we
        # still want to know which tryjobs exist.
        config_name = config_name or "commit-queue.cfg"
        return self.fetch_config(config_name, cv_config_pb2.Config, **kwargs)

    def milo(self, **kwargs):
        return self.fetch_config("luci-milo.cfg", milo_pb2.Project, **kwargs)

    def get_milo_console(self, console_name, **kwargs):
        milo_cfg = self.milo(**kwargs)
        for console in milo_cfg.consoles:
            if console.id == console_name:
                return console
        return None