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.cq.api.config.v2 import cq as cq_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):
        """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.
        """
        if not project:
            project = self.m.buildbucket.build.builder.project
            assert project, "buildbucket input has no project set"

        with self.m.step.nest("fetch %s" % config_name):
            endpoint = "config/v1/config_sets"
            config_set = "projects/%s/config" % project
            url = "https://luci-config.appspot.com/_ah/api/{}/{}/{}".format(
                endpoint, config_set, config_name
            )
            token = self.m.service_account.default().get_access_token()
            response = self.m.url.get_json(
                url,
                step_name="get",
                headers={"Authorization": "Bearer {}".format(token)},
            )

        text = base64.b64decode(response.output["content"])
        cfg = message_type()
        textpb.Parse(text, cfg)
        return cfg

    def commit_queue(self, **kwargs):
        return self.fetch_config("commit-queue.cfg", cq_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