[Invoke] Feature: Add required checks using invoke instead of Github …

…UI (DataDog#31243)

tasks/github_tasks.py

@@ -18,7 +18,7 @@
     print_workflow_conclusion,
     trigger_macos_workflow,
 )
-from tasks.libs.common.color import color_message
+from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.constants import DEFAULT_INTEGRATIONS_CORE_BRANCH
 from tasks.libs.common.datadog_api import create_gauge, send_event, send_metrics
 from tasks.libs.common.git import get_default_branch
@@ -546,3 +546,36 @@ def agenttelemetry_list_change_ack_check(_, pr_id=-1):
             print(
                 "'need-change/agenttelemetry-governance' label found on the PR: potential change to Agent Telemetry metrics is acknowledged and the governance instructions are followed."
             )
+
+
+@task
+def get_required_checks(_, branch: str = "main"):
+    """
+    For this task to work:
+        - A Personal Access Token (PAT) needs the "repo" permissions.
+        - A fine-grained token needs the "Administration" repository permissions (read).
+    """
+    from tasks.libs.ciproviders.github_api import GithubAPI
+
+    gh = GithubAPI()
+    required_checks = gh.get_branch_required_checks(branch)
+    print(required_checks)
+
+
+@task(iterable=['check'])
+def add_required_checks(_, branch: str, check: str, force: bool = False):
+    """
+    For this task to work:
+        - A Personal Access Token (PAT) needs the "repo" permissions.
+        - A fine-grained token needs the "Administration" repository permissions (write).
+
+    Use it like this:
+    inv github.add-required-checks --branch=main --check="dd-gitlab/lint_codeowners" --check="dd-gitlab/lint_components"
+    """
+    from tasks.libs.ciproviders.github_api import GithubAPI
+
+    if not check:
+        raise Exit(color_message("No check name provided, exiting", Color.RED), code=1)
+
+    gh = GithubAPI()
+    gh.add_branch_required_check(branch, check, force)

tasks/libs/ciproviders/github_api.py

@@ -11,9 +11,10 @@
 
 import requests
 
-from tasks.libs.common.color import color_message
+from tasks.libs.common.color import Color, color_message
 from tasks.libs.common.constants import GITHUB_REPO_NAME
 from tasks.libs.common.git import get_default_branch
+from tasks.libs.common.user_interactions import yes_no_question
 
 try:
     import semver
@@ -90,6 +91,136 @@ def get_milestone_by_name(self, milestone_name):
                 return milestone
         return None
 
+    def get_branch_protection(self, branch_name: str):
+        """
+        Get the protection of a given branch
+        """
+        branch = self.get_branch(branch_name)
+        if not branch:
+            raise Exit(color_message(f"Branch {branch_name} not found", Color.RED), code=1)
+        elif not branch.protected:
+            raise Exit(color_message(f"Branch {branch_name} doesn't have protection", Color.RED), code=1)
+        try:
+            protection = branch.get_protection()
+        except GithubException as e:
+            if e.status == 403:
+                error_msg = f"""Can't access {branch_name} branch protection, probably due to missing permissions. You need either:
+    - A Personal Access Token (PAT) needs the "repo" permissions.
+    - Or a fine-grained token needs the "Administration" repository permissions.
+"""
+                raise PermissionError(error_msg) from e
+            raise
+        return protection
+
+    def protection_to_payload(self, protection_raw_data: dict) -> dict:
+        """
+        Convert the protection object to a payload.
+        See https://docs.github.com/en/rest/branches/branch-protection?apiVersion=2022-11-28#update-branch-protection
+
+        The following seems to be defined at the Org scale, so we're not resending them here:
+        - required_pull_request_reviews > dismissal_restrictions
+        - required_pull_request_reviews > bypass_pull_request_allowances
+        """
+        prot = protection_raw_data
+        return {
+            "required_status_checks": {
+                "strict": prot["required_status_checks"]["strict"],
+                "checks": [
+                    {"context": check["context"], "app_id": -1 if check["app_id"] is None else check["app_id"]}
+                    for check in prot["required_status_checks"]["checks"]
+                ],
+            },
+            "enforce_admins": prot["enforce_admins"]["enabled"],
+            "required_pull_request_reviews": {
+                "dismiss_stale_reviews": prot["required_pull_request_reviews"]["dismiss_stale_reviews"],
+                "require_code_owner_reviews": prot["required_pull_request_reviews"]["require_code_owner_reviews"],
+                "required_approving_review_count": prot["required_pull_request_reviews"][
+                    "required_approving_review_count"
+                ],
+                "require_last_push_approval": prot["required_pull_request_reviews"]["require_last_push_approval"],
+            },
+            "restrictions": {
+                "users": prot["restrictions"]["users"],
+                "teams": prot["restrictions"]["teams"],
+                "apps": [app["slug"] for app in prot["restrictions"]["apps"]],
+            },
+            "required_linear_history": prot["required_linear_history"]["enabled"],
+            "allow_force_pushes": prot["allow_force_pushes"]["enabled"],
+            "allow_deletions": prot["allow_deletions"]["enabled"],
+            "block_creations": prot["block_creations"]["enabled"],
+            "required_conversation_resolution": prot["required_conversation_resolution"]["enabled"],
+            "lock_branch": prot["lock_branch"]["enabled"],
+            "allow_fork_syncing": prot["allow_fork_syncing"]["enabled"],
+        }
+
+    def get_branch_required_checks(self, branch_name: str) -> list[str]:
+        """
+        Get the required checks for a given branch
+        """
+        return self.get_branch_protection(branch_name).required_status_checks.contexts
+
+    def add_branch_required_check(self, branch_name: str, checks: list[str], force: bool = False) -> None:
+        """
+        Add required checks to a given branch
+
+        It uses the Github API directly to add the required checks to the branch.
+        Using the "checks" argument is not supported by PyGithub.
+        :calls: `PUT /repos/{owner}/{repo}/branches/{branch}/protection
+
+        """
+        current_protection = self.get_branch_protection(branch_name)
+        current_required_checks = current_protection.required_status_checks.contexts
+        new_required_checks = []
+        for check in checks:
+            if check in current_required_checks:
+                print(
+                    color_message(
+                        f"Ignoring the '{check}' check as it is already required on the {branch_name} branch",
+                        Color.ORANGE,
+                    )
+                )
+            else:
+                new_required_checks.append(check)
+        if not new_required_checks:
+            print(color_message("No new checks to add", Color.GREEN))
+            return
+        print(
+            color_message(
+                f"Warning: You are about to add the following checks to the {branch_name} branch:\n{new_required_checks}",
+                Color.ORANGE,
+            )
+        )
+        print(color_message(f"Current required checks: {sorted(current_required_checks)}", Color.GREY))
+        if force or yes_no_question("Are you sure?", default=False):
+            # We're crafting the request and not using PyGithub because it doesn't support passing the checks variable instead of contexts.
+            protection_url = f"{self.repo.url}/branches/{branch_name}/protection"
+            headers = {
+                "Accept": "application/vnd.github+json",
+                "Authorization": f"Bearer {os.environ['GITHUB_TOKEN']}",
+                "X-GitHub-Api-Version": "2022-11-28",
+            }
+            payload = self.protection_to_payload(current_protection.raw_data)
+            payload["required_status_checks"]["checks"] = sorted(
+                payload["required_status_checks"]["checks"]
+                + [{"context": check, "app_id": -1} for check in new_required_checks],
+                key=lambda x: x['context'],
+            )
+
+            response = requests.put(protection_url, headers=headers, json=payload, timeout=10)
+            if response.status_code != 200:
+                print(
+                    color_message(
+                        f"Error while sending the PUT request to {protection_url}\n{response.text}", Color.RED
+                    )
+                )
+                raise Exit(
+                    color_message(f"Failed to update the required checks for the {branch_name} branch", Color.RED),
+                    code=1,
+                )
+            print(color_message(f"The {checks} checks were successfully added!", Color.GREEN))
+        else:
+            print(color_message("Aborting changes to the branch required checks", Color.GREEN))
+
     def is_release_note_needed(self, pull_number):
         """
         Check if labels are ok for skipping QA