datastore/ci/praktika/_environment.py

import dataclasses
import json
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Dict, List, Optional, Type

from . import Job, Workflow
from .settings import Settings
from .utils import MetaClasses, Shell, T, Utils


@dataclasses.dataclass
class _Environment(MetaClasses.Serializable):
    WORKFLOW_NAME: str
    JOB_NAME: str
    REPOSITORY: str
    BRANCH: str
    SHA: str
    PR_NUMBER: int
    EVENT_TYPE: str
    EVENT_TIME: str
    JOB_OUTPUT_STREAM: str
    EVENT_FILE_PATH: str
    CHANGE_URL: str
    COMMIT_URL: str
    BASE_BRANCH: str
    RUN_ID: str
    RUN_URL: str
    INSTANCE_TYPE: str
    INSTANCE_ID: str
    INSTANCE_LIFE_CYCLE: str
    PR_BODY: str
    PR_TITLE: str
    USER_LOGIN: str
    FORK_NAME: str
    COMMIT_MESSAGE: str = ""
    # merged PR for "push" or "merge_group" workflow
    LINKED_PR_NUMBER: int = 0
    LOCAL_RUN: bool = False
    PR_LABELS: List[str] = dataclasses.field(default_factory=list)
    REPORT_MESSAGES: List[Dict[str, str]] = dataclasses.field(default_factory=list)
    JOB_CONFIG: Optional[Job.Config] = None
    TRACEBACKS: List[str] = dataclasses.field(default_factory=list)
    WORKFLOW_JOB_DATA: Dict[str, Any] = dataclasses.field(default_factory=dict)
    JOB_KV_DATA: Dict[str, Any] = dataclasses.field(default_factory=dict)
    # Keys present in JOB_KV_DATA when the job started, i.e. inherited from the
    # initial (config) job. Used in Runner._post_run to emit only the KV data a
    # job itself added as the job's `data` output, instead of re-emitting the
    # whole inherited bucket into every job's output (toJson(needs)).
    JOB_KV_DATA_BASE_KEYS: List[str] = dataclasses.field(default_factory=list)
    COMMIT_AUTHORS: List[str] = dataclasses.field(default_factory=list)
    WORKFLOW_CONFIG: Optional[Dict[str, Any]] = None
    name = "environment"

    @classmethod
    def _load_workflow_job_data(cls) -> dict:
        if Path(Settings.WORKFLOW_JOB_FILE).is_file():
            with open(Settings.WORKFLOW_JOB_FILE, "r", encoding="utf8") as f:
                return json.load(f)
        print(
            f"NOTE: Workflow job file [{Settings.WORKFLOW_JOB_FILE}] does not exist"
        )
        return {}

    @classmethod
    def from_env(cls) -> "_Environment":
        WORKFLOW_NAME = os.getenv("GITHUB_WORKFLOW", "")
        JOB_NAME = os.getenv("JOB_NAME", "")
        REPOSITORY = os.getenv("GITHUB_REPOSITORY", "")
        # GITHUB_HEAD_REF for pull_request, GITHUB_REF_NAME for push
        BRANCH = os.getenv("GITHUB_HEAD_REF", "") or os.getenv("GITHUB_REF_NAME", "")

        EVENT_FILE_PATH = os.getenv("GITHUB_EVENT_PATH", "")
        JOB_OUTPUT_STREAM = os.getenv("GITHUB_OUTPUT", "")
        RUN_ID = os.getenv("GITHUB_RUN_ID", "0")
        RUN_URL = f"https://github.com/{REPOSITORY}/actions/runs/{RUN_ID}"
        BASE_BRANCH = os.getenv("GITHUB_BASE_REF", "")
        USER_LOGIN = ""
        COMMIT_AUTHORS = []
        FORK_NAME = REPOSITORY
        PR_BODY = ""
        PR_TITLE = ""
        PR_LABELS = []
        LINKED_PR_NUMBER = 0
        EVENT_TIME = ""
        COMMIT_MESSAGE = ""

        WORKFLOW_JOB_DATA = cls._load_workflow_job_data()

        if EVENT_FILE_PATH:
            with open(EVENT_FILE_PATH, "r", encoding="utf-8") as f:
                github_event = json.load(f)
            if "pull_request" in github_event:
                FORK_NAME = github_event["pull_request"]["head"]["repo"]["full_name"]
                EVENT_TYPE = Workflow.Event.PULL_REQUEST
                PR_NUMBER = int(github_event["pull_request"]["number"])
                SHA = github_event["pull_request"]["head"]["sha"]
                CHANGE_URL = github_event["pull_request"]["html_url"]
                COMMIT_URL = CHANGE_URL + f"/commits/{SHA}"
                PR_BODY = github_event["pull_request"]["body"]
                PR_TITLE = github_event["pull_request"]["title"]
                PR_LABELS = [
                    label["name"] for label in github_event["pull_request"]["labels"]
                ]
                USER_LOGIN = github_event["pull_request"]["user"]["login"]
                EVENT_TIME = github_event.get("pull_request", {}).get(
                    "updated_at", None
                )
            elif "commits" in github_event:
                EVENT_TYPE = Workflow.Event.PUSH
                SHA = github_event["after"]
                CHANGE_URL = github_event["head_commit"]["url"]  # commit url
                PR_NUMBER = 0
                COMMIT_URL = CHANGE_URL
                commit_message = github_event["head_commit"]["message"]
                parts = commit_message.split(
                    "pull request #"
                )  # Merge pull request #77106 from XXXX
                if len(parts) >= 2:
                    pr_str = parts[1].split()[0]
                    LINKED_PR_NUMBER = int(pr_str) if pr_str.isdigit() else 0
                EVENT_TIME = github_event.get("repository", {}).get("updated_at", None)
                commit_authors = set()
                for commit in github_event["commits"]:
                    email = commit["author"]["email"]
                    # Validate email contains @ symbol
                    if email and "@" in email:
                        commit_authors.add(email)
                COMMIT_AUTHORS = list(commit_authors)
            elif "schedule" in github_event:
                EVENT_TYPE = Workflow.Event.SCHEDULE
                SHA = os.getenv(
                    "GITHUB_SHA", "0000000000000000000000000000000000000000"
                )
                PR_NUMBER = 0
                CHANGE_URL = (
                    github_event["repository"]["html_url"] + "/commit/" + SHA
                )  # commit url
                COMMIT_URL = CHANGE_URL
            elif "inputs" in github_event:
                # assume this is a dispatch
                EVENT_TYPE = Workflow.Event.DISPATCH
                SHA = os.getenv(
                    "GITHUB_SHA", "0000000000000000000000000000000000000000"
                )
                PR_NUMBER = 0
                CHANGE_URL = (
                    github_event["repository"]["html_url"] + "/commit/" + SHA
                )  # commit url
                COMMIT_URL = CHANGE_URL
            elif "merge_group" in github_event:
                PR_NUMBER = 0
                EVENT_TYPE = Workflow.Event.MERGE_QUEUE
                SHA = os.getenv(
                    "GITHUB_SHA", "0000000000000000000000000000000000000000"
                )
                head_ref = github_event.get("merge_group", {}).get("head_ref", "")
                try:
                    if "/pr-" in head_ref:
                        pr_number_part = head_ref.split("/pr-")[1].split("-")[0]
                        LINKED_PR_NUMBER = (
                            int(pr_number_part) if pr_number_part.isdigit() else 0
                        )
                    CHANGE_URL = (
                        f"{github_event['repository']['html_url']}/pull/{LINKED_PR_NUMBER}"
                        if LINKED_PR_NUMBER
                        else ""
                    )
                except:
                    LINKED_PR_NUMBER = 0
                    CHANGE_URL = ""

                COMMIT_URL = CHANGE_URL
            else:
                assert False, "TODO: not supported"

            INSTANCE_TYPE = (
                os.getenv("INSTANCE_TYPE", None)
                or Shell.get_output("ec2metadata --instance-type")
                or ""
            )
            INSTANCE_ID = (
                os.getenv("INSTANCE_ID", None)
                or Shell.get_output("ec2metadata --instance-id")
                or ""
            )
            INSTANCE_LIFE_CYCLE = (
                os.getenv("INSTANCE_LIFE_CYCLE", None)
                or Shell.get_output(
                    "curl -s --fail http://169.254.169.254/latest/meta-data/instance-life-cycle"
                )
                or ""
            )

        else:
            print("WARNING: Local execution - dummy Environment will be generated")
            SHA = "TEST"
            PR_NUMBER = -1
            EVENT_TYPE = Workflow.Event.PUSH
            CHANGE_URL = ""
            COMMIT_URL = ""
            COMMIT_MESSAGE = ""
            INSTANCE_TYPE = ""
            INSTANCE_ID = ""
            INSTANCE_LIFE_CYCLE = ""

        return _Environment(
            WORKFLOW_NAME=WORKFLOW_NAME,
            JOB_NAME=JOB_NAME,
            REPOSITORY=REPOSITORY,
            BRANCH=BRANCH,
            EVENT_FILE_PATH=EVENT_FILE_PATH,
            JOB_OUTPUT_STREAM=JOB_OUTPUT_STREAM,
            SHA=SHA,
            EVENT_TYPE=EVENT_TYPE,
            EVENT_TIME=EVENT_TIME,
            PR_NUMBER=PR_NUMBER,
            RUN_ID=RUN_ID,
            CHANGE_URL=CHANGE_URL,
            COMMIT_URL=COMMIT_URL,
            RUN_URL=RUN_URL,
            BASE_BRANCH=BASE_BRANCH,
            INSTANCE_TYPE=INSTANCE_TYPE,
            INSTANCE_ID=INSTANCE_ID,
            PR_BODY=PR_BODY,
            PR_TITLE=PR_TITLE,
            USER_LOGIN=USER_LOGIN,
            COMMIT_AUTHORS=COMMIT_AUTHORS,
            FORK_NAME=FORK_NAME,
            COMMIT_MESSAGE=COMMIT_MESSAGE,
            PR_LABELS=PR_LABELS,
            INSTANCE_LIFE_CYCLE=INSTANCE_LIFE_CYCLE,
            REPORT_MESSAGES=[],
            LINKED_PR_NUMBER=LINKED_PR_NUMBER,
            # TODO: Find a better way to store and pass commit authors data through workflow
            JOB_KV_DATA={
                "commit_authors": COMMIT_AUTHORS,
                # Initial parent PR inference:
                # - Defaults to LINKED_PR_NUMBER, which is the PR merged by a push/merge-queue event
                # - Can be explicitly overridden later via workflow hooks (see Info.set_parent_pr_number)
                "parent_pr_number": LINKED_PR_NUMBER,
            },
            WORKFLOW_JOB_DATA=WORKFLOW_JOB_DATA,
            WORKFLOW_CONFIG=None,
        )

    @classmethod
    def file_name_static(cls, _name=""):
        return f"{Settings.TEMP_DIR}/{cls.name}.json"

    @classmethod
    def from_dict(cls: Type[T], obj: Dict[str, Any]) -> T:
        JOB_OUTPUT_STREAM = os.getenv("GITHUB_OUTPUT", "")
        obj["JOB_OUTPUT_STREAM"] = JOB_OUTPUT_STREAM
        if "PARAMETER" in obj:
            obj["PARAMETER"] = _to_object(obj["PARAMETER"])
        # Filter out unexpected arguments - only keep fields defined in the dataclass
        valid_fields = {f.name for f in dataclasses.fields(cls)}
        filtered_obj = {k: v for k, v in obj.items() if k in valid_fields}
        return cls(**filtered_obj)

    @classmethod
    def from_workflow_data(cls) -> "_Environment":
        with open(Settings.WORKFLOW_STATUS_FILE, "r", encoding="utf8") as f:
            workflow_status_data = json.load(f)
        # Access the config job data and parse the JSON string in "data" field
        # Job names are normalized in the workflow status file
        normalized_job_name = Utils.normalize_string(Settings.CI_CONFIG_JOB_NAME)
        config_job_data = workflow_status_data.get(normalized_job_name, {})
        data_str = config_job_data.get("outputs", {}).get("data", "")
        # A suppressed output may surface either as a missing key or as a blank
        # string - do not let json.loads choke on the latter before the
        # explicit error below
        if isinstance(data_str, str):
            env_dict = json.loads(data_str) if data_str.strip() else None
        else:
            env_dict = data_str
        if not env_dict or "SHA" not in env_dict:
            raise RuntimeError(
                f"Job output [data] of the [{Settings.CI_CONFIG_JOB_NAME}] job is empty or missing. "
                "If that job itself succeeded, the runner has likely suppressed the output: "
                "it refuses to export a job output that matches a secret pattern "
                "(look for a [##[warning]Skip output 'data' since it may contain secret] line "
                f"at the very end of the [{Settings.CI_CONFIG_JOB_NAME}] job log)"
            )

        # JOB_KV_DATA is stored as opaque base64 in the initial job's output
        # (see Runner.run) to keep user-authored strings such as changed file
        # paths from matching a secret pattern and suppressing the output -
        # decode it back into a dict here
        kv_data = env_dict.get("JOB_KV_DATA")
        if isinstance(kv_data, str):
            env_dict["JOB_KV_DATA"] = (
                json.loads(Utils.from_base64(kv_data)) if kv_data else {}
            )

        # User-authored free text is stripped from the initial job's output
        # (see Runner.run: the runner drops outputs matching secret patterns) -
        # restore it from the local event payload
        event_file_path = os.getenv("GITHUB_EVENT_PATH", "")
        if event_file_path and Path(event_file_path).is_file():
            with open(event_file_path, "r", encoding="utf-8") as f:
                github_event = json.load(f)
            if "pull_request" in github_event:
                env_dict["PR_BODY"] = github_event["pull_request"]["body"] or ""
                env_dict["PR_TITLE"] = github_event["pull_request"]["title"] or ""
            elif github_event.get("head_commit"):
                env_dict["COMMIT_MESSAGE"] = (
                    github_event["head_commit"]["message"] or ""
                )

        # Reread instance metadata from the host
        env_dict["INSTANCE_TYPE"] = (
            os.getenv("INSTANCE_TYPE", None)
            or Shell.get_output("ec2metadata --instance-type")
            or ""
        )
        env_dict["INSTANCE_ID"] = (
            os.getenv("INSTANCE_ID", None)
            or Shell.get_output("ec2metadata --instance-id")
            or ""
        )
        env_dict["INSTANCE_LIFE_CYCLE"] = (
            os.getenv("INSTANCE_LIFE_CYCLE", None)
            or Shell.get_output(
                "curl -s --fail http://169.254.169.254/latest/meta-data/instance-life-cycle"
            )
            or ""
        )

        # Override WORKFLOW_JOB_DATA with the current job's data so that
        # check_run_id refers to this job rather than to the config job whose
        # serialised environment we loaded above.
        env_dict["WORKFLOW_JOB_DATA"] = cls._load_workflow_job_data()

        return cls.from_dict(env_dict)

    def _add_report_message(self, message, kind, source=""):
        """
        Accumulate a structured report message in the environment.

        Messages are collected during job execution and later written to both
        the job and workflow ``Result.ext`` as ``{"message": str, "from": str}``
        entries.  Grouping of duplicate messages is done at the rendering level
        in ``json.html``.
        Prefer the typed wrappers ``add_workflow_warning/error/note``.
        """
        self.REPORT_MESSAGES.append(
            {"message": message, "kind": kind, "from": source or self.JOB_NAME}
        )
        self.dump()

    def add_workflow_warning(self, message, source=""):
        self._add_report_message(message, kind="warning", source=source)

    def add_workflow_error(self, message, source=""):
        self._add_report_message(message, kind="error", source=source)

    def add_workflow_note(self, message, source=""):
        self._add_report_message(message, kind="note", source=source)

    @classmethod
    def get(cls):
        if Path(cls.file_name_static()).is_file():
            return cls.from_fs("environment")
        else:
            try:
                env = cls.from_workflow_data()
                env.dump()
                return env
            except FileNotFoundError:
                # For workflows without Config job
                print(
                    f"NOTE: Workflow context file [{Settings.WORKFLOW_STATUS_FILE}] does not exist - read context from GH event"
                )
                env = cls.from_env()
                env.dump()
                return env

    def set_job_name(self, job_name):
        self.JOB_NAME = job_name
        self.dump()
        return self

    def set_pr_labels(self, labels: List[str], reset: bool = False) -> None:
        if reset:
            self.PR_LABELS = list(labels)
        else:
            for label in labels:
                if label not in self.PR_LABELS:
                    self.PR_LABELS.append(label)
        self.dump()

    def add_pr_label(self, label: str) -> None:
        if label not in self.PR_LABELS:
            self.PR_LABELS.append(label)
            self.dump()

    def remove_pr_label(self, label: str) -> None:
        if label in self.PR_LABELS:
            self.PR_LABELS.remove(label)
            self.dump()

    @staticmethod
    def get_needs_statuses():
        if Path(Settings.WORKFLOW_STATUS_FILE).is_file():
            with open(Settings.WORKFLOW_STATUS_FILE, "r", encoding="utf8") as f:
                return json.load(f)
        else:
            print(
                f"ERROR: Status file [{Settings.WORKFLOW_STATUS_FILE}] does not exist"
            )
            raise RuntimeError()

    def get_s3_prefix(self, latest=False):
        return self.get_s3_prefix_static(self.PR_NUMBER, self.BRANCH, self.SHA, latest)

    @classmethod
    def get_s3_prefix_static(cls, pr_number, branch, sha, latest=False):
        assert pr_number > 0 or branch
        if pr_number:
            prefix = f"PRs/{pr_number}"
        else:
            prefix = f"REFs/{branch}"
        assert sha or latest
        if latest:
            prefix += "/latest"
        elif sha:
            prefix += f"/{sha}"
        return prefix

    def is_local_run(self):
        return self.LOCAL_RUN


def _to_object(data):
    if isinstance(data, dict):
        return SimpleNamespace(**{k: _to_object(v) for k, v in data.items()})
    elif isinstance(data, list):
        return [_to_object(i) for i in data]
    else:
        return data