bin/ci/ci_post_gantt.py

#!/usr/bin/env python3
# Copyright © 2023 Collabora Ltd.
# Authors:
#   Helen Koike <helen.koike@collabora.com>
#
# For the dependencies, see the requirements.txt
# SPDX-License-Identifier: MIT


import argparse
import logging as log
import os
import re
import traceback
from datetime import datetime, timedelta
from typing import Any, Dict

import gitlab
import pytz
from ci_gantt_chart import generate_gantt_chart
from gitlab import Gitlab
from gitlab.base import RESTObject
from gitlab.v4.objects import Project, ProjectPipeline
from gitlab_common import (GITLAB_URL, get_gitlab_pipeline_from_url,
                           get_token_from_default_dir, read_token)


class MockGanttExit(Exception):
    pass


LAST_MARGE_EVENT_FILE = os.path.expanduser("~/.config/last_marge_event")


def read_last_event_date_from_file() -> str:
    try:
        with open(LAST_MARGE_EVENT_FILE, "r") as f:
            last_event_date = f.read().strip()
    except FileNotFoundError:
        # 3 days ago
        last_event_date = (datetime.now() - timedelta(days=3)).isoformat()
    return last_event_date


def pretty_time(time_str: str) -> str:
    """Pretty print time"""
    local_timezone = datetime.now().astimezone().tzinfo

    time_d = datetime.fromisoformat(time_str.replace("Z", "+00:00")).astimezone(
        local_timezone
    )
    return f'{time_str} ({time_d.strftime("%d %b %Y %Hh%Mm%Ss")} {local_timezone})'


def compose_message(file_name: str, attachment_url: str) -> str:
    return f"""
[{file_name}]({attachment_url})

<details>
<summary>more info</summary>

This message was generated by the ci_post_gantt.py script, which is running on a server at Collabora.
</details>
"""


def gitlab_upload_file_get_url(gl: Gitlab, project_id: str, filepath: str) -> str:
    project: Project = gl.projects.get(project_id)
    uploaded_file: Dict[str, Any] = project.upload(filepath, filepath=filepath)
    return uploaded_file["url"]


def gitlab_post_reply_to_note(gl: Gitlab, event: RESTObject, reply_message: str):
    """
    Post a reply to a note in thread based on a GitLab event.

    :param gl: The GitLab connection instance.
    :param event: The event object containing the note details.
    :param reply_message: The reply message.
    """
    try:
        note_id = event.target_id
        merge_request_iid = event.note["noteable_iid"]

        project = gl.projects.get(event.project_id)
        merge_request = project.mergerequests.get(merge_request_iid)

        # Find the discussion to which the note belongs
        discussions = merge_request.discussions.list(iterator=True)
        target_discussion = next(
            (
                d
                for d in discussions
                if any(n["id"] == note_id for n in d.attributes["notes"])
            ),
            None,
        )

        if target_discussion is None:
            raise ValueError("Discussion for the note not found.")

        # Add a reply to the discussion
        reply = target_discussion.notes.create({"body": reply_message})
        return reply

    except gitlab.exceptions.GitlabError as e:
        log.error(f"Failed to post a reply to '{event.note['body']}': {e}")
        return None


def main(
    token: str | None,
    since: str | None,
    marge_user_id: int = 9716,
    project_ids: list[int] = [176],
    ci_timeout: float = 60,
):
    log.basicConfig(level=log.INFO)
    if token is None:
        token = get_token_from_default_dir()

    token = read_token(token)
    gl = Gitlab(url=GITLAB_URL, private_token=token, retry_transient_errors=True)

    user = gl.users.get(marge_user_id)
    last_event_at = since if since else read_last_event_date_from_file()

    log.info(f"Retrieving Marge messages since {pretty_time(last_event_at)}\n")

    # the "after" only considers the "2023-10-24" part, it doesn't consider the time
    events = user.events.list(
        all=True,
        target_type="note",
        after=(datetime.now() - timedelta(days=3)).isoformat(),
        sort="asc",
    )

    last_event_at_date = datetime.fromisoformat(
        last_event_at.replace("Z", "+00:00")
    ).replace(tzinfo=pytz.UTC)

    for event in events:
        if event.project_id not in project_ids:
            continue
        created_at_date = datetime.fromisoformat(
            event.created_at.replace("Z", "+00:00")
        ).replace(tzinfo=pytz.UTC)
        if created_at_date <= last_event_at_date:
            continue
        last_event_at = event.created_at

        escaped_gitlab_url = re.escape(GITLAB_URL)
        match = re.search(rf"{escaped_gitlab_url}/[^\s<]+", event.note["body"])

        if match:
            try:
                log.info(f"Found message: {event.note['body']}")
                pipeline_url = match.group(0)[:-1]
                pipeline: ProjectPipeline
                pipeline, _ = get_gitlab_pipeline_from_url(gl, pipeline_url)
                log.info("Generating gantt chart...")
                fig = generate_gantt_chart(pipeline, ci_timeout)
                file_name = f"{str(pipeline.id)}-Gantt.html"
                fig.write_html(file_name)
                log.info("Uploading gantt file...")
                file_url = gitlab_upload_file_get_url(gl, event.project_id, file_name)
                log.info("Posting reply ...")
                message = compose_message(file_name, file_url)
                gitlab_post_reply_to_note(gl, event, message)
            except MockGanttExit:
                pass  # Allow tests to exit early without printing a traceback
            except Exception as e:
                log.info(f"Failed to generate gantt chart, not posting reply.{e}")
                traceback.print_exc()

        if not since:
            log.info(
                f"Updating last event date to {pretty_time(last_event_at)} on {LAST_MARGE_EVENT_FILE}\n"
            )
            with open(LAST_MARGE_EVENT_FILE, "w") as f:
                f.write(last_event_at)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Monitor rejected pipelines by Marge.")
    parser.add_argument(
        "--token",
        metavar="token",
        type=str,
        default=None,
        help="force GitLab token, otherwise it's read from ~/.config/gitlab-token",
    )
    parser.add_argument(
        "--since",
        metavar="since",
        type=str,
        default=None,
        help="consider only events after this date (ISO format), otherwise it's read from ~/.config/last_marge_event",
    )
    parser.add_argument(
        "--marge-user-id",
        metavar="marge_user_id",
        type=int,
        default=9716,  # default https://gitlab.freedesktop.org/users/marge-bot/activity
        help="GitLab user ID for marge-bot, defaults to 9716",
    )
    parser.add_argument(
        "--project-id",
        metavar="project_id",
        type=int,
        nargs="+",
        default=[176],  # default is the mesa/mesa project id
        help="GitLab project id(s) to analyze. Defaults to 176 i.e. mesa/mesa.",
    )
    parser.add_argument(
        "--ci-timeout",
        metavar="ci_timeout",
        type=float,
        default=60,
        help="Time that marge-bot will wait for ci to finish. Defaults to one hour.",
    )
    args = parser.parse_args()
    main(args.token, args.since, args.marge_user_id, args.project_id, args.ci_timeout)