ci-status-overview

#!/usr/bin/env python3

'''Show an overview of what CI checks are failing or succeeding.

The overview links to check results in supporting terminals.
'''

from __future__ import annotations
import abc
import argparse
import glob
import io
import itertools as it
import math
import os
import os.path
import shutil
import sys
from collections import defaultdict
from collections.abc import Iterable, Mapping   # for type checking
from datetime import datetime, timedelta
from textwrap import dedent
from typing import Optional, Union, Literal, Final, TypedDict, TextIO

from gql import Client, gql
from gql.transport.requests import RequestsHTTPTransport
from graphql.language.ast import DocumentNode
from alibot_helpers.utilities import parse_env_file


DEFAULTENV: Final[str] = 'DEFAULTS.env'
TIMEFORMAT: Final[str] = '%Y-%m-%dT%H:%M:%SZ'
NOW: Final[datetime] = datetime.now()
PR_URL_FMT: Final[str] = 'https://github.com/{repo}/pull/{pr}'
API_URL: Final[str] = 'https://api.github.com/graphql'
GET_PR_STATUSES_GRAPHQL: Final[DocumentNode] = gql('''\
query statuses($repoOwner: String!, $repoName: String!, $baseBranch: String!) {
  repository(owner: $repoOwner, name: $repoName) {
    pullRequests(last: 100, baseRefName: $baseBranch, states: OPEN) {
      nodes {
        number
        title
        isDraft
        commits(last: 1) {
          nodes {
            commit {
              oid
              status {
                contexts {
                  context
                  state
                  createdAt
                  targetUrl
                }
              }
            }
          }
        }
      }
    }
  }
}
''')


State = Union[Literal['EXPECTED'], Literal['ERROR'], Literal['FAILURE'],
              Literal['PENDING'], Literal['SUCCESS']]
VALID_STATUSES: Final[tuple[State, ...]] = \
    'EXPECTED', 'PENDING', 'FAILURE', 'ERROR', 'SUCCESS'


class Check(TypedDict, total=False):
    '''The result of a single check of a commit.'''
    state: State
    context: str
    createdAt: str
    repo: str
    pr: int
    commit_sha: str
    targetUrl: str


def get_status_url(status: Check) -> Optional[str]:
    '''Construct a useful URL for the given check, falling back to its PR.'''
    url = status.get('targetUrl')
    if url:
        return url
    try:
        return PR_URL_FMT.format(**status)
    except KeyError:
        return None


def get_check_statuses(client, repo: str, branch: str, checks: list[str]) \
        -> Mapping[str, list[Check]]:
    '''Yield {check: [status]} for all given checks on PRs in repo.'''
    owner, is_valid, repo_name = repo.partition('/')
    if not is_valid:
        raise ValueError('repository name must contain a slash')
    statuses: defaultdict[str, list[Check]] = defaultdict(list)
    response = client.execute(GET_PR_STATUSES_GRAPHQL, {
        'repoOwner': owner, 'repoName': repo_name, 'baseBranch': branch,
    })
    for pull in response['repository']['pullRequests']['nodes']:
        if pull['isDraft'] or pull['title'].startswith('[WIP]'):
            continue
        # We only ever get one commit in the response from GitHub.
        commit = pull['commits']['nodes'][0]['commit']
        contexts = ({c['context']: c for c in commit['status']['contexts']}
                    if commit['status'] else {})
        for check in checks:
            # Fallback to 'expected' status with sensible defaults.
            fallback = {'context': check, 'state': 'EXPECTED',
                        'createdAt': NOW.strftime(TIMEFORMAT)}
            statuses[check].append(contexts.get(check, fallback) | {
                'repo': repo,
                'pr': pull['number'],
                'commit_sha': commit['oid'],
            })
    return statuses


def get_all_checks(defs_dir: str, roles: list[str], containers: list[str],
                   repos: list[str], checks: list[str]) \
        -> Mapping[tuple[str, str], list[str]]:
    '''Parse .env files and return checks in each repo.'''
    all_checks: defaultdict[tuple[str, str], list[str]] = defaultdict(list)
    for env_path in glob.glob(os.path.join(defs_dir, '*', '*', '*.env')):
        if env_path.endswith(os.sep + DEFAULTENV):
            continue
        check = {}
        role, docker, env = os.path.relpath(env_path, defs_dir).split(os.sep)
        if roles and role not in roles:
            continue
        if containers and docker not in containers:
            continue
        for envpath in (os.path.join(defs_dir, DEFAULTENV),
                        os.path.join(defs_dir, role, DEFAULTENV),
                        os.path.join(defs_dir, role, docker, DEFAULTENV),
                        os.path.join(defs_dir, role, docker, env)):
            if os.path.exists(envpath):
                check.update(parse_env_file(envpath))
        repo, branch, name = \
            check['PR_REPO'], check['PR_BRANCH'], check['CHECK_NAME']
        if repos and repo not in repos:
            continue
        if checks and name not in checks:
            continue
        all_checks[(repo, branch)].append(name)
    return all_checks


class Output(abc.ABC):
    '''The base class for output formatters.'''

    def __init__(self: Output, recent_hours: float,
                 output_file: TextIO = sys.stdout) -> None:
        self.recent_hours = recent_hours
        self.recent_cutoff = (NOW - timedelta(hours=recent_hours)).strftime(TIMEFORMAT)
        self.output_file = output_file

    def document(self: Output, client,
                 all_checks: Mapping[tuple[str, str], list[str]]) -> None:
        '''Display the entire overview document for the given checks.'''
        self.begin()
        for (repo, branch), checks in sorted(all_checks.items()):
            self.repo_header(repo, branch)
            statuses = get_check_statuses(client, repo, branch, checks)
            for check in sorted(checks):
                self.check_header(check)
                if statuses[check]:
                    self.overview_table(statuses[check])
                else:
                    self.empty_table()
        self.end()

    def begin(self: Output) -> None:
        '''Output anything that is required to be before the overview table.'''

    @abc.abstractmethod
    def repo_header(self: Output, repo: str, branch: str) -> None:
        '''Display the repository and branch name.'''

    @abc.abstractmethod
    def check_header(self: Output, check_name: str) -> None:
        '''Display the given check name.'''

    @abc.abstractmethod
    def empty_table(self: Output) -> None:
        '''Display a helpful message for a table with no PRs.'''

    @staticmethod
    def overview_table_prep(statuses: list[Check]) -> tuple[int, str]:
        '''Sort statuses and return the length and template of labels.'''
        statuses.sort(key=lambda c: c['createdAt'], reverse=True)
        prnum_len = 1 + math.floor(math.log10(max(c['pr'] for c in statuses)))
        return prnum_len, f'#{{:{prnum_len}d}}'

    @abc.abstractmethod
    def overview_table(self: Output, pr_statuses: list[Check]) -> None:
        '''Show a nicely formatted table of PR results for the given check.'''

    def end(self: Output) -> None:
        '''Output anything that is required to be after the overview table.'''


class TextOutput(Output):
    '''Show the overview on the terminal in nice colours.'''
    INDENT: Final[str] = '  '
    SEPARATOR: Final[str] = '  '

    def repo_header(self: TextOutput, repo: str, branch: str) -> None:
        '''Format repo bold and underlined, and italicize branch.'''
        print(f'\033[4;1m{repo}\033[0m  \033[3m({branch})\033[0m',
              file=self.output_file)

    def check_header(self: TextOutput, check_name: str) -> None:
        '''Format the given check name by underlining.'''
        print(f'{self.INDENT}\033[4m{check_name}\033[0m',
              file=self.output_file)

    def empty_table(self: TextOutput) -> None:
        '''Create a helpful message for a table with no PRs.'''
        print(self.INDENT, self.INDENT,
              '\033[3;90m(no open non-draft PRs here)\033[0m',
              sep='', end='\n\n', file=self.output_file)

    def overview_table(self: TextOutput, pr_statuses: list[Check]) -> None:
        '''Print a nicely formatted table of PR results for the given check.'''
        prnum_len, template = self.overview_table_prep(pr_statuses)
        terminal_width = shutil.get_terminal_size((80, 1)).columns
        items_per_row = ((terminal_width - 2*len(self.INDENT) + len(self.SEPARATOR)) //
                         (len('#') + prnum_len + len(self.SEPARATOR)))
        for _, row in it.groupby(enumerate(pr_statuses),
                                 key=lambda tpl: tpl[0] // items_per_row):
            self.table_row(((status, template.format(status['pr']))
                            for _, status in row))
        print(file=self.output_file)

    def table_row(self: TextOutput,
                  statuses_and_text: Iterable[tuple[Check, str]]) -> None:
        '''Format each text for its accompanying status and optional URL.'''
        print(self.INDENT, self.INDENT, sep='', end='', file=self.output_file)
        print(*(self.format_status(status, text)
                for status, text in statuses_and_text),
              sep=self.SEPARATOR, file=self.output_file)

    def format_status(self: TextOutput, status: Check, text: str) -> str:
        '''Color the given text as appropriate for the given status.

        If possible, the text is also formatted as a hyperlink to the document
        reporting check results.
        '''
        ansi_escaped = ''
        url = get_status_url(status)
        if url:
            ansi_escaped += '\033]8;;' + url + '\033\\'  # opening URL code

        ansi_escaped += '\033[' + {   # start opening color code
            'PENDING': '33',    # yellow
            'EXPECTED': '90',   # gray (bright black)
            'SUCCESS': '32',    # green
            'ERROR': '31',      # red
            'FAILURE': '31;1',  # bold red
        }[status['state']]

        if status['createdAt'] > self.recent_cutoff:
            ansi_escaped += ';7'  # reverse video -- swap fore- and background

        ansi_escaped += 'm' + text   # finish color code and append text
        if url:
            ansi_escaped += '\033]8;;\033\\'  # closing URL code
        return ansi_escaped + '\033[0m'  # closing color code


class HtmlOutput(Output):
    '''Export the overview table as a HTML document.'''

    def begin(self: HtmlOutput) -> None:
        '''Output the HTML <head> element and initial boilerplate.'''
        print(dedent('''\
        <!doctype html>
        <html>
        <head>
        <meta charset="utf-8">
        <title>ALICE CI overview</title>
        <style type="text/css">
        body { font-family: sans-serif; margin: 0; padding: 1rem; }
        #key { padding: 0.5rem; }
        #key[open] { border: 0.15rem dashed #777; }
        #key summary { font-weight: bold; }
        .branch-name { font-family: monospace; font-size: 1.25rem;
                       font-style: italic; margin-left: 0.75rem; }
        .branch-name::before { content: "("; }
        .branch-name::after { content: ")"; }
        .check-name { margin-left: 1rem; }
        .empty { font-size: 0.875rem; color: #777; font-style: italic; }
        .table { margin-left: 1.75rem; display: flex; place-content: start;
                 flex-flow: row wrap; }
        .status { padding: 0.25rem; margin: 0.25rem; --status-color: currentColor;
                  border: 0.1rem solid transparent; color: var(--status-color); }
        .status a { display: block; color: inherit; }
        .status.recent { border-color: var(--status-color); }
        .status.EXPECTED { --status-color: #24292f; border-style: dotted; }
        .status.PENDING { --status-color: #bf8700; }
        .status.SUCCESS { --status-color: #1a7f37; }
        .status.ERROR { --status-color: #cf222e; }
        .status.FAILURE { --status-color: #cf222e; font-weight: bold; }
        </style>
        </head>
        <body>
        <h1>ALICE CI overview</h1>
        '''), dedent(f'''\
        <p>Document generated at {NOW.strftime(TIMEFORMAT)}. Statuses from the
           last <strong>{self.recent_hours:g} hours</strong>, i.e. newer than
           {self.recent_cutoff}, are marked as
           <span class="status recent">recent</span>.</p>
        '''), dedent('''\
        <details id="key"><summary>Explanation (click to expand)</summary>
        <p>The results of the check listed in each heading are shown for each
           pull request in a list.</p>
        <p>Results are ordered most recent first.</p>
        <p>Checks that completed after a set cutoff point (see the top of this
           document for the specific time) have a border around them,
           <span class="status recent">like this</span>.</p>
        <p>The colour coding works as follows:</p>
        <ul>
          <li><span class="status EXPECTED">#0000</span> is an "expected"
              status, which means that the CI has not picked up this PR at all
              yet for the respective check.</li>
          <li><span class="status PENDING">#0000</span> is a "pending" status,
              which means that the CI has picked this PR up, but the check has
              not yet completed.</li>
          <li><span class="status SUCCESS">#0000</span> is a successful status,
              i.e. this check has run and no errors were found.</li>
          <!--<li><span class="status FAILURE">#0000</span> is a failed status,
              which doesn't happen with the current CI system.</li>-->
          <li><span class="status ERROR">#0000</span> is an error status, which
              means that the check has run but a build error occurred.</li>
        </ul>
        </details>
        '''), sep='', end='', file=self.output_file)

    def repo_header(self: HtmlOutput, repo: str, branch: str) -> None:
        '''Output a heading with the repository and branch name.'''
        print(f'<h2>{repo} <span class="branch-name">{branch}</span></h2>',
              file=self.output_file)

    def check_header(self: HtmlOutput, check_name: str) -> None:
        '''Output the given check name.'''
        print(f'<h3 class="check-name">{check_name}</h3>',
              file=self.output_file)

    def empty_table(self: HtmlOutput) -> None:
        '''Output a helpful message for a table with no PRs.'''
        print('<div class="table empty">(no open non-draft PRs here)</div>',
              file=self.output_file)

    def overview_table(self: HtmlOutput, pr_statuses: list[Check]) -> None:
        '''Show a nicely formatted table of PR results for the given check.'''
        _, template = self.overview_table_prep(pr_statuses)
        print('<div class="table">',
              *(self.format_status(status, template.format(status['pr']))
                for status in pr_statuses),
              '</div>', sep='\n', file=self.output_file)

    def format_status(self: HtmlOutput, status: Check, text: str) -> str:
        '''Tag the given text as appropriate for the given status.

        If possible, the text is also hyperlinked to the document reporting
        check results.
        '''
        if (url := get_status_url(status)):
            text = f'<a href="{url}">{text}</a>'
        date = 'recent' if status['createdAt'] > self.recent_cutoff else 'old'
        return (f'<div class="status {status["state"]} {date}"'
                f' title="{status.get("createdAt", "")}">{text}</div>')

    def end(self: HtmlOutput) -> None:
        '''Close any open tags.'''
        print('</body></html>', file=self.output_file)

def parse_args() -> argparse.Namespace:
    '''Parse and return command-line args.'''
    with io.StringIO() as demo_io:
        demo_text_output = TextOutput(recent_hours=1, output_file=demo_io)
        demo_text_output.repo_header('owner/repository', 'base branch')
        demo_text_output.check_header('check name')
        demo_text_output.table_row((
            ({'state': status, 'createdAt': NOW.strftime(TIMEFORMAT)},
             status.lower() + ' (recent)')
            for status in VALID_STATUSES
        ))
        demo_text_output.table_row((
            ({'state': status,
              'createdAt': (NOW - timedelta(hours=2)).strftime(TIMEFORMAT)},
             status.lower() + ' (older) ')
            for status in VALID_STATUSES
        ))
        table = demo_io.getvalue()

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=dedent('''\
        This script needs a GITHUB_TOKEN environment variable.

        Key to the output:
        {table}

        "Expected" means no builder has picked the check up yet. Draft PRs
        (those marked as drafts and those with "[WIP]" in the title) are not
        shown. Read tables left to right, then down.
        ''').format(table=table))

    parser.add_argument(
        '-o', '--html-output', dest='html_output_file', metavar='FILE',
        type=argparse.FileType('wt'),
        help=('output a complete HTML document containing the overview '
              'to %(metavar)s (use "-" for stdout)'))
    parser.add_argument(
        '--definitions-dir', metavar='DIR', default='ali-bot/ci/repo-config',
        help=('directory where .env files are located in a hierarchy; expects '
              'a directory structure of the form %(metavar)s/ROLE/CONTAINER/'
              '*.env (default %(default)s)'))
    parser.add_argument(
        '-t', '--recent-hours', metavar='HOURS', type=float, default=24.0,
        help=('consider check results from the last %(metavar)s days recent '
              '(these are printed in reverse video; default %(default)g; '
              'this can be a non-integer value)'))

    filtering = parser.add_argument_group('filter displayed checks', dedent('''\
    Each filtering argument can be given multiple times. If no filtering
    arguments are given, all known checks are shown in the output.

    Filtering arguments can be combined. If multiple are given (possibly
    multiple times each), the criteria are OR-ed together. For example, "-c
    check1 -r repo1 -r repo2" would show an overview of check1 (in any repo),
    in addition to all checks in repo1 or repo2.
    '''))
    filtering.add_argument(
        '-m', '--mesos-role', action='append', metavar='ROLE', dest='roles',
        default=[], help='include checks running under this Mesos role')
    filtering.add_argument(
        '-d', '--docker-container', action='append', metavar='CONTAINER',
        dest='containers', default=[],
        help=('include checks running inside this Docker container (use the '
              'short name only, e.g. alisw/slc8-builder:latest -> slc8)'))
    filtering.add_argument(
        '-r', '--repo', action='append', metavar='USER/REPO', dest='repos',
        default=[],
        help=('include checks for this repository (of the form '
              "<user>/<repository>; don't include github.com)"))
    filtering.add_argument(
        '-c', '--check', action='append', metavar='NAME', dest='checks',
        default=[],
        help=('include the specific named check (use the name as it appears '
              'on GitHub, e.g. build/O2/o2)'))

    args = parser.parse_args()
    if 'GITHUB_TOKEN' not in os.environ:
        parser.error('GITHUB_TOKEN environment variable is required')
    return args


def main(args: argparse.Namespace) -> None:
    '''Main entry point.'''
    output: Output
    if args.html_output_file:
        output = HtmlOutput(args.recent_hours, args.html_output_file)
    else:
        output = TextOutput(args.recent_hours)
    with Client(transport=RequestsHTTPTransport(url=API_URL, headers={
            'Authorization': 'bearer ' + os.environ['GITHUB_TOKEN'],
    }), fetch_schema_from_transport=False) as session:
        output.document(session, get_all_checks(
            args.definitions_dir, args.roles, args.containers, args.repos,
            args.checks))


if __name__ == '__main__':
    main(parse_args())