[SPARK-9383][PROJECT-INFRA] PR merge script should reset back to previous branch when possible

This patch modifies our PR merge script to reset back to a named branch when restoring the original checkout upon exit. When the committer is originally checked out to a detached head, then they will be restored back to that same ref (the same as today's behavior).

This is a slightly updated version of #7569, with an extra fix to handle the detached head corner-case.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #10709 from JoshRosen/SPARK-9383.

1 files changed, 16 insertions, 3 deletions

diff --git a/dev/merge_spark_pr.py b/dev/merge_spark_pr.py
index bf1a000f46..5ab285eae9 100755
--- a/dev/merge_spark_pr.py
+++ b/dev/merge_spark_pr.py
@@ -355,11 +355,21 @@ def standardize_jira_ref(text):
 
     return clean_text
 
+
+def get_current_ref():
+    ref = run_cmd("git rev-parse --abbrev-ref HEAD").strip()
+    if ref == 'HEAD':
+        # The current ref is a detached HEAD, so grab its SHA.
+        return run_cmd("git rev-parse HEAD").strip()
+    else:
+        return ref
+
+
 def main():
     global original_head
 
     os.chdir(SPARK_HOME)
-    original_head = run_cmd("git rev-parse HEAD")[:8]
+    original_head = get_current_ref()
 
     branches = get_json("%s/branches" % GITHUB_API_BASE)
     branch_names = filter(lambda x: x.startswith("branch-"), [x['name'] for x in branches])
@@ -449,5 +459,8 @@ if __name__ == "__main__":
     (failure_count, test_count) = doctest.testmod()
     if failure_count:
         exit(-1)
-
-    main()
+    try:
+        main()
+    except:
+        clean_up()
+        raise