[SPARK-57002][INFRA] Enforce Upstream-First policy in merge_spark_pr.py cherry-pick prompts

dev/merge_spark_pr.py

@@ -470,11 +470,8 @@ def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, pr_author, co_author
     return merge_hash
 
 
-def cherry_pick(pr_num, merge_hash, default_branch):
-    pick_ref = bold_input("Enter a branch name [%s]: " % default_branch)
-    if pick_ref == "":
-        pick_ref = default_branch
-
+def _do_cherry_pick(pr_num, merge_hash, pick_ref):
+    """Cherry-pick `merge_hash` onto `pick_ref` and push. Returns the pushed ref."""
     pick_branch_name = "%s_PICK_PR_%s_%s" % (BRANCH_PREFIX, pr_num, pick_ref.upper())
 
     run_cmd("git fetch %s %s:%s" % (PUSH_REMOTE_NAME, pick_ref, pick_branch_name))
@@ -495,7 +492,6 @@ def cherry_pick(pr_num, merge_hash, default_branch):
     try:
         run_cmd("git push %s %s:%s" % (PUSH_REMOTE_NAME, pick_branch_name, pick_ref))
     except Exception as e:
-        clean_up()
         fail("Exception while pushing: %s" % e)
 
     pick_hash = run_cmd("git rev-parse %s" % pick_branch_name)[:8]
@@ -506,6 +502,84 @@ def cherry_pick(pr_num, merge_hash, default_branch):
     return pick_ref
 
 
+def _upstream_first_sibling(target_ref, pick_ref, branch_names, already_picked):
+    """Return the sibling branch-M.x if Upstream-First should prompt, else None.
+
+    The policy only applies when the PR was merged into master: that's the only case
+    where the committer can type branch-M.N at the cherry-pick prompt and bypass the
+    rolling branch-M.x. When the PR was opened against branch-M.x the merge itself
+    lands there (nothing to bypass), and when it was opened against branch-M.N the
+    author already chose per-branch scope.
+
+    >>> _upstream_first_sibling("master", "branch-4.2", ["branch-4.x", "branch-4.2"], ())
+    'branch-4.x'
+    >>> _upstream_first_sibling("master", "branch-4.2", ["branch-4.x", "branch-4.2"],
+    ...                         ("branch-4.x",))
+    >>> _upstream_first_sibling("master", "branch-4.x", ["branch-4.x"], ())
+    >>> _upstream_first_sibling("master", "branch-4.99", ["branch-4.2"], ())
+    >>> _upstream_first_sibling("branch-4.x", "branch-4.2", ["branch-4.x", "branch-4.2"], ())
+    >>> _upstream_first_sibling("branch-4.2", "branch-3.5", ["branch-4.x", "branch-3.5"], ())
+    """
+    if target_ref != "master":
+        return None
+    m = re.match(r"^branch-(\d+)\.(\d+)$", pick_ref)
+    if not m:
+        return None
+    candidate = "branch-%s.x" % m.group(1)
+    if candidate in branch_names and candidate not in already_picked:
+        return candidate
+    return None
+
+
+def cherry_pick(pr_num, merge_hash, default_branch, branch_names, target_ref, already_picked=()):
+    """Prompt for a target branch and cherry-pick `merge_hash` onto it.
+
+    Enforces the Upstream-First policy (see header comment) via
+    `_upstream_first_sibling`: when the PR was merged into master and the committer
+    types a branch-M.N target while branch-M.x is also a known release branch AND
+    has not already received this commit, prompt to confirm whether to pick into
+    BOTH (the policy-compliant default) or branch-M.N only (treated as a
+    maintenance-only bugfix). Returns the list of refs actually picked into, so
+    the main loop can advance its remaining-branches list correctly.
+    """
+    pick_ref = bold_input("Enter a branch name [%s]: " % default_branch)
+    if pick_ref == "":
+        pick_ref = default_branch
+
+    sibling_x = _upstream_first_sibling(target_ref, pick_ref, branch_names, already_picked)
+    if sibling_x is not None:
+        print()
+        print("=" * 80)
+        print(
+            "Upstream-First policy: non-bugfix commits on %s should also land on %s."
+            % (pick_ref, sibling_x)
+        )
+        print("If this is a %s-only maintenance bugfix, you may pick %s alone."
+              % (pick_ref, pick_ref))
+        print("Otherwise, pick both (%s first, then %s)." % (sibling_x, pick_ref))
+        print("=" * 80)
+        choice = (
+            bold_input(
+                "Pick into [b]oth %s + %s / [o]nly %s / [a]bort (default: both): "
+                % (sibling_x, pick_ref, pick_ref)
+            )
+            .strip()
+            .lower()
+        )
+        if choice in ("", "b", "both"):
+            picked_x = _do_cherry_pick(pr_num, merge_hash, sibling_x)
+            picked_n = _do_cherry_pick(pr_num, merge_hash, pick_ref)
+            return [picked_x, picked_n]
+        elif choice in ("o", "only"):
+            return [_do_cherry_pick(pr_num, merge_hash, pick_ref)]
+        elif choice in ("a", "abort"):
+            fail("Aborted by user at Upstream-First policy prompt.")
+        else:
+            fail("Unrecognized choice %r; aborting." % choice)
+
+    return [_do_cherry_pick(pr_num, merge_hash, pick_ref)]
+
+
 def print_jira_issue_summary(issue):
     summary = "Summary\t\t%s\n" % issue.fields.summary
     assignee = issue.fields.assignee
@@ -832,7 +906,6 @@ def main():
     branches = get_json("%s/branches" % GITHUB_API_BASE)
     branch_names = list(filter(lambda x: x.startswith("branch-"), [x["name"] for x in branches]))
     branch_names = sorted(branch_names, key=semver_branch_rank, reverse=True)
-    branch_iter = iter(branch_names)
 
     if len(sys.argv) == 1:
         pr_num = bold_input("Which pull request would you like to merge? (e.g. 34): ")
@@ -942,7 +1015,8 @@ def main():
             fail("Couldn't find any merge commit for #%s, you may need to update HEAD." % pr_num)
 
         print("Found commit %s:\n%s" % (merge_hash, message))
-        cherry_pick(pr_num, merge_hash, next(branch_iter, branch_names[0]))
+        default = branch_names[0]
+        cherry_pick(pr_num, merge_hash, default, branch_names, target_ref, already_picked=())
         sys.exit(0)
 
     if not bool(pr["mergeable"]):
@@ -976,11 +1050,23 @@ def main():
         print("PR #%s is still open after push; closing it explicitly." % pr_num)
         close_pr(pr_num)
 
+    # Walk a mutable remaining-branches list so the next default correctly skips any
+    # branches already picked, including branches consumed by the Upstream-First two-branch
+    # path inside cherry_pick (e.g. picking branch-M.x + branch-M.N in a single prompt).
+    # merged_refs doubles as the already_picked set passed to cherry_pick: it starts with
+    # target_ref (the merge sink, never to be re-picked) and grows with every cherry-pick.
+    remaining_branches = [b for b in branch_names if b != target_ref]
     pick_prompt = "Would you like to pick %s into another branch?" % merge_hash
     while bold_input("\n%s (y/N): " % pick_prompt).lower() == "y":
-        merged_refs = merged_refs + [
-            cherry_pick(pr_num, merge_hash, next(branch_iter, branch_names[0]))
-        ]
+        default = remaining_branches[0] if remaining_branches else branch_names[0]
+        picked = cherry_pick(
+            pr_num, merge_hash, default, branch_names, target_ref,
+            already_picked=tuple(merged_refs),
+        )
+        merged_refs = merged_refs + picked
+        for b in picked:
+            if b in remaining_branches:
+                remaining_branches.remove(b)
 
     if asf_jira is not None:
         continue_maybe("Would you like to update an associated JIRA?")