mesa/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 gitlab
import re
import os
import pytz
import traceback
from datetime import datetime, timedelta
from gitlab_common import (
    read_token,
    GITLAB_URL,
    get_gitlab_pipeline_from_url,
)
from ci_gantt_chart import generate_gantt_chart

MARGE_USER_ID = 9716  # Marge

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


def read_last_event_date_from_file():
    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):
    """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, attachment_url):
    return f"""
Here is the Gantt chart for the referred pipeline, I hope it helps 😄 (tip: click on the "Pan" button on the top right bar):

[{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, project_id, filepath):
    project = gl.projects.get(project_id)
    uploaded_file = project.upload(filepath, filepath=filepath)
    return uploaded_file["url"]


def gitlab_post_reply_to_note(gl, event, reply_message):
    """
    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(as_list=False)
        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:
        print(f"Failed to post a reply to '{event.note['body']}': {e}")
        return None


def parse_args() -> None:
    parser = argparse.ArgumentParser(description="Monitor rejected pipelines by Marge.")
    parser.add_argument(
        "--token",
        metavar="token",
        help="force GitLab token, otherwise it's read from ~/.config/gitlab-token",
    )
    parser.add_argument(
        "--since",
        metavar="since",
        help="consider only events after this date (ISO format), otherwise it's read from ~/.config/last_marge_event",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    token = read_token(args.token)

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

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

    print(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:
        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

        match = re.search(r"https://[^ ]+", event.note["body"])
        if match:
            try:
                print("Found message:", event.note["body"])
                pipeline_url = match.group(0)[:-1]
                pipeline, _ = get_gitlab_pipeline_from_url(gl, pipeline_url)
                print("Generating gantt chart...")
                fig = generate_gantt_chart(pipeline)
                file_name = "Gantt.html"
                fig.write_html(file_name)
                print("Uploading gantt file...")
                file_url = gitlab_upload_file_get_url(gl, event.project_id, file_name)
                print("Posting reply ...\n")
                message = compose_message(file_name, file_url)
                gitlab_post_reply_to_note(gl, event, message)
            except Exception as e:
                print(f"Failed to generate gantt chart, not posting reply.{e}")
                traceback.print_exc()

        if not args.since:
            print(
                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)