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.
TREE_STATUS_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f"


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)
    username = attr.ib(type=str)
    # Unique ID of this status.
    key = attr.ib(type=int)

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

    @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, 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.
        """
        step = self.m.step(
            step_name,
            [self._tree_status_script, hostname, "get"],
            stdout=self.m.json.output(),
            step_test_data=self.test_api.status_step_data,
        )
        status = TreeStatus.from_dict(**step.stdout)
        step.presentation.step_text = status.state.value
        step.presentation.links[hostname] = "https://%s" % hostname
        return status

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

        Args:
            hostname (str): Hostname of the tree status page.
            password (str): Password for the service account to access 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.
            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, step_name="check for tree status collision"
            )
            if current_status.key != last_status.key:
                raise self.m.step.StepFailure(
                    "mid-air collision detected between tree status changes"
                )
        return self.m.step(
            step_name,
            [
                self._tree_status_script,
                hostname,
                "set",
                message,
                "--username",
                username or self.m.buildbucket.builder_name,
                "--password",
                password,
            ],
        )