recipe_modules/tree_status/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 attr
import datetime
import enum

from recipe_engine import recipe_api

# The format used for the "date" field in HTTP responses from requests to the
# tree status page.
OLD_TREE_STATUS_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f"
TREE_STATUS_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"


class State(enum.Enum):
    OPEN = "OPEN"
    CLOSED = "CLOSED"
    THROTTLED = "THROTTLED"
    MAINTENANCE = "MAINTENANCE"


@attr.s
class TreeStatus:
    state = attr.ib(type=str)
    date = attr.ib(type=datetime.datetime)
    # Unique ID of this status.
    key = attr.ib(type=str)
    message = attr.ib(type=str)
    username = attr.ib(type=str, default=None)

    @classmethod
    def from_dict(cls, name, generalState, createTime, message, **_):
        """Create a TreeStatus from a dict, ignoring any unneeded values."""
        return cls(
            state=State(generalState),
            date=datetime.datetime.strptime(createTime, TREE_STATUS_DATE_FORMAT),
            message=message,
            # The `name` will be in the form of `trees/{tree_name}/status/{key}`.
            key=name.split("/")[-1].strip(),
        )

    @classmethod
    def old_status_from_dict(cls, general_state, date, username, key, message, **_):
        """Create a TreeStatus from a dict, ignoring any unneeded values."""
        return cls(
            state=State(general_state.upper()),
            date=datetime.datetime.strptime(date, OLD_TREE_STATUS_DATE_FORMAT),
            username=username,
            key=str(key),
            message=message,
        )

    @property
    def open(self):
        return self.state in (State.OPEN, State.THROTTLED)


class TreeStatusApi(recipe_api.RecipeApi):
    @property
    def _tree_status_script(self):
        return self.resource("tree_status.py")

    def get(self, hostname, tree_name=None, step_name="get current tree status"):
        """Get the current tree status.

        Args:
            hostname (str): Hostname of the tree status page.

        Returns:
            A `TreeStatus` corresponding to the current state of the tree.
        """
        cmd = [
            self._tree_status_script,
            hostname,
        ]
        if tree_name:
            cmd.extend(["--tree-name", tree_name, "--use-new-status-app"])
        cmd.append("get")
        step = self.m.step(
            step_name,
            cmd,
            stdout=self.m.json.output(),
            step_test_data=self.test_api.status_step_data,
        )
        if tree_name:
            status = TreeStatus.from_dict(**step.stdout)
        else:
            status = TreeStatus.old_status_from_dict(**step.stdout)
        step.presentation.step_text = status.state.value
        if tree_name:
            step.presentation.links[tree_name] = (
                f"https://ci.chromium.org/ui/labs/tree-status/{tree_name}"
            )
        else:
            step.presentation.links[hostname] = f"https://{hostname}"
        return status

    def update(
        self,
        hostname,
        admin_hostname,
        message,
        tree_name=None,
        username=None,
        state=None,
        issuetracker_id=None,
        last_status=None,
        step_name="update tree status",
    ):
        """Change the tree status text.

        Args:
            hostname (str): Hostname of the tree status page.
            message (str): Message to set as the tree status.
            username (str): Name or email to use as the author of the tree
                status.
            state (str): The state of the tree as a single word (e.g. "CLOSED", "OPEN")
            issuetracker_id (str): The ID of the filed IssueTracker bug.
            last_status (TreeStatus or None): Only update the tree status if it
                has not changed since the specified status change.
        """
        # Collision checking *should* be handled by the tree status app itself,
        # but it doesn't do collision checking for status changes that use the
        # bot-specific endpoint (as opposed to the web UI form):
        # https://chromium.googlesource.com/infra/infra/+/09d53b324dd786b96d69c6cbb8ee6c389b22fd3f/appengine/chromium_status/appengine_module/chromium_status/status.py#426
        if last_status:
            current_status = self.get(
                hostname,
                tree_name=tree_name,
                step_name="check for tree status collision",
            )
            if current_status.key != last_status.key:
                raise self.m.step.StepFailure(
                    "collision detected between tree status changes"
                )
        cmd = [
            self._tree_status_script,
            hostname,
        ]
        if tree_name:
            cmd.extend(["--tree-name", tree_name, "--use-new-status-app"])
        cmd.extend(
            [
                "set",
                message,
                "--admin_hostname",
                admin_hostname,
                "--username",
                username or self.m.buildbucket.builder_name,
            ]
        )
        if state:
            cmd.extend(["--state", state])
        if issuetracker_id:
            cmd.extend(["--issuetracker_id", issuetracker_id])
        return self.m.step(step_name, cmd)