[Utils] Fixed rebase in merge-release-pr script (#116340)

Recently GitHub changed something on their side so we no longer can
rebase release PR's with the API. This means that we now have to
manually rebase the PR locally and then push the results. This fixes
the script that I use to merge PRs to the release branch by changing
the rebase part to do the local rebase and also adds a new option
--rebase-only so that you can rebase the PRs easier.

Minor change is that the script now can take a URL to the pull request
as well as just the PR ID.

Author: tru

llvm/utils/release/merge-release-pr.py

@@ -106,27 +106,61 @@ def validate_commits(self, data):
             return False, f"More than 1 commit! {len(data['commits'])}"
         return True
 
-    def validate_pr(self):
+    def _normalize_pr(self, parg: str):
+        if parg.isdigit():
+            return parg
+        elif parg.startswith("https://github.com/llvm/llvm-project/pull"):
+            # try to parse the following url https://github.com/llvm/llvm-project/pull/114089
+            i = parg[parg.rfind("/") + 1 :]
+            if not i.isdigit():
+                raise RuntimeError(f"{i} is not a number, malformatted input.")
+            return i
+        else:
+            raise RuntimeError(
+                f"PR argument must be PR ID or pull request URL - {parg} is wrong."
+            )
+
+    def load_pr_data(self):
+        self.args.pr = self._normalize_pr(self.args.pr)
         fields_to_fetch = [
             "baseRefName",
+            "commits",
+            "headRefName",
+            "headRepository",
+            "headRepositoryOwner",
             "reviewDecision",
-            "title",
+            "state",
             "statusCheckRollup",
+            "title",
             "url",
-            "state",
-            "commits",
         ]
+        print(f"> Loading PR {self.args.pr}...")
         o = self.run_gh(
             "pr",
             ["view", self.args.pr, "--json", ",".join(fields_to_fetch)],
         )
-        prdata = json.loads(o)
+        self.prdata = json.loads(o)
 
         # save the baseRefName (target branch) so that we know where to push
-        self.target_branch = prdata["baseRefName"]
+        self.target_branch = self.prdata["baseRefName"]
+        srepo = self.prdata["headRepository"]["name"]
+        sowner = self.prdata["headRepositoryOwner"]["login"]
+        self.source_url = f"https://github.com/{sowner}/{srepo}"
+        self.source_branch = self.prdata["headRefName"]
+
+        if srepo != "llvm-project":
+            print("The target repo is NOT llvm-project, check the PR!")
+            sys.exit(1)
+
+        if sowner == "llvm":
+            print(
+                "The source owner should never be github.com/llvm, double check the PR!"
+            )
+            sys.exit(1)
 
-        print(f"> Handling PR {self.args.pr} - {prdata['title']}")
-        print(f">   {prdata['url']}")
+    def validate_pr(self):
+        print(f"> Handling PR {self.args.pr} - {self.prdata['title']}")
+        print(f">   {self.prdata['url']}")
 
         VALIDATIONS = {
             "state": self.validate_state,
@@ -141,7 +175,7 @@ def validate_pr(self):
         total_ok = True
         for val_name, val_func in VALIDATIONS.items():
             try:
-                validation_data = val_func(prdata)
+                validation_data = val_func(self.prdata)
             except:
                 validation_data = False
             ok = None
@@ -166,24 +200,42 @@ def validate_pr(self):
         return total_ok
 
     def rebase_pr(self):
-        print("> Rebasing")
-        self.run_gh("pr", ["update-branch", "--rebase", self.args.pr])
-        print("> Waiting for GitHub to update PR")
-        time.sleep(4)
+        print("> Fetching upstream")
+        subprocess.run(["git", "fetch", "--all"], check=True)
+        print("> Rebasing...")
+        subprocess.run(
+            ["git", "rebase", self.args.upstream + "/" + self.target_branch], check=True
+        )
+        print("> Publish rebase...")
+        subprocess.run(
+            ["git", "push", "--force", self.source_url, f"HEAD:{self.source_branch}"]
+        )
 
     def checkout_pr(self):
         print("> Fetching PR changes...")
+        self.merge_branch = "llvm_merger_" + self.args.pr
         self.run_gh(
             "pr",
             [
                 "checkout",
                 self.args.pr,
                 "--force",
                 "--branch",
-                "llvm_merger_" + self.args.pr,
+                self.merge_branch,
             ],
         )
 
+        # get the branch information so that we can use it for
+        # pushing later.
+        p = subprocess.run(
+            ["git", "config", f"branch.{self.merge_branch}.merge"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        upstream_branch = p.stdout.strip().replace("refs/heads/", "")
+        print(upstream_branch)
+
     def push_upstream(self):
         print("> Pushing changes...")
         subprocess.run(
@@ -201,7 +253,7 @@ def delete_local_branch(self):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "pr",
-        help="The Pull Request ID that should be merged into a release.",
+        help="The Pull Request ID that should be merged into a release. Can be number or URL",
     )
     parser.add_argument(
         "--skip-validation",
@@ -224,9 +276,20 @@ def delete_local_branch(self):
     parser.add_argument(
         "--validate-only", action="store_true", help="Only run the validations."
     )
+    parser.add_argument(
+        "--rebase-only", action="store_true", help="Only rebase and exit"
+    )
     args = parser.parse_args()
 
     merger = PRMerger(args)
+    merger.load_pr_data()
+
+    if args.rebase_only:
+        merger.checkout_pr()
+        merger.rebase_pr()
+        merger.delete_local_branch()
+        sys.exit(0)
+
     if not merger.validate_pr():
         print()
         print(
@@ -239,8 +302,8 @@ def delete_local_branch(self):
         print("! --validate-only passed, will exit here")
         sys.exit(0)
 
-    merger.rebase_pr()
     merger.checkout_pr()
+    merger.rebase_pr()
 
     if args.no_push:
         print()