From 1deb4017f5021fb5a269a483e582cf6eaef298bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9lian=20Raimbault?=
 <161456554+CelianR@users.noreply.github.com>
Date: Thu, 9 Jan 2025 08:15:46 -0500
Subject: [PATCH] Add task to print pr state (#32806)

---
 tasks/github_tasks.py                | 30 ++++++++++++++++++++++++++++
 tasks/libs/ciproviders/github_api.py | 16 +++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/tasks/github_tasks.py b/tasks/github_tasks.py
index a71e00b834997..50696bb3d197d 100644
--- a/tasks/github_tasks.py
+++ b/tasks/github_tasks.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import os
 import re
 import time
@@ -639,3 +640,32 @@ def check_qa_labels(_, labels: str):
     if len(qa_labels) > 1:
         raise Exit(f"More than one QA label set.\n{docs}", code=1)
     print("QA label set correctly")
+
+
+@task
+def print_pr_state(_, id):
+    """Print the PR merge state if the PR is stuck within the merge queue."""
+
+    from tasks.libs.ciproviders.github_api import GithubAPI
+
+    query = """
+query {
+  repository (owner: "DataDog", name: "datadog-agent") {
+    pullRequest(number: ID) {
+      reviewDecision
+      state
+      statusCheckRollup {
+        state
+      }
+      mergeable
+      mergeStateStatus
+      locked
+    }
+  }
+}
+""".replace("ID", id)  # Use replace to avoid formatting issues with curly braces
+
+    gh = GithubAPI()
+    res = gh.graphql(query)
+
+    print(json.dumps(res, indent=2))
diff --git a/tasks/libs/ciproviders/github_api.py b/tasks/libs/ciproviders/github_api.py
index 63fcfeb09a204..fc0850b79ec4a 100644
--- a/tasks/libs/ciproviders/github_api.py
+++ b/tasks/libs/ciproviders/github_api.py
@@ -53,6 +53,22 @@ def repo(self):
 
         return self._repository
 
+    def graphql(self, query):
+        """
+        Perform a GraphQL query against the Github API.
+        """
+
+        headers = {"Authorization": "Bearer " + self._auth.token, "Content-Type": "application/json"}
+        res = requests.post(
+            "https://api.github.com/graphql",
+            headers=headers,
+            json={"query": query},
+        )
+        if res.status_code == 200:
+            return res.json()
+        else:
+            raise RuntimeError(f"Failed to query Github: {res.text}")
+
     def get_branch(self, branch_name):
         """
         Gets info on a given branch in the given Github repository.