[AIRFLOW-176] Improve PR Tool JIRA workflow

- Fix crash when non-integer IDs are passed
- Improve workflow by always asking user if they
  want to resolve another issue before exiting

dev/airflow-pr

@@ -296,7 +296,55 @@ def fix_version_from_branch(branch, versions):
             return versions[-1]
 
 
+def validate_jira_id(jira_id):
+    if not jira_id:
+        return
+
+    # first look for AIRFLOW-X
+    ids = re.findall("AIRFLOW-[0-9]{1,6}", jira_id)
+    if len(ids) > 1:
+        raise click.UsageError('Found multiple issue ids: {}'.format(ids))
+    elif len(ids) == 1:
+        jira_id = ids[0]
+    elif not ids:
+        # if we don't find AIRFLOW-X, see if jira_id is an int
+        try:
+            jira_id = 'AIRFLOW-{}'.format(abs(int(jira_id)))
+        except ValueError:
+            raise click.UsageError(
+                'JIRA id must be an integer or have the form AIRFLOW-X')
+
+    return jira_id
+
+
+def resolve_jira_issues_loop(comment=None, merge_branches=None):
+    """
+    Resolves a JIRA issue, then asks the user if he/she would like to close
+    another one. Repeats until the user indicates they are finished.
+    """
+    while True:
+        try:
+            resolve_jira_issue(
+                comment=comment,
+                jira_id=None,
+                merge_branches=merge_branches)
+        except Exception as e:
+            click.echo("ERROR: {}".format(e))
+
+        if not click.confirm('Would you like to resolve another JIRA issue?'):
+            return
+
+
 def resolve_jira_issue(comment=None, jira_id=None, merge_branches=None):
+    """
+    Resolves a JIRA issue
+
+    comment: a comment for the issue. The user will always be prompted for one;
+        if provided, this will be the default.
+
+    jira_id: an Airflow JIRA id, either an integer or a string with the form
+        AIRFLOW-X. If not provided, the user will be prompted to provide one.
+    """
     if merge_branches is None:
         merge_branches = []
 
@@ -314,20 +362,17 @@ def resolve_jira_issue(comment=None, jira_id=None, merge_branches=None):
         {'server': JIRA_API_BASE},
         basic_auth=(JIRA_USERNAME, JIRA_PASSWORD))
 
-    jira_id = 'AIRFLOW-{}'.format(abs(click.prompt(
-        'Enter an Airflow JIRA id', default=jira_id, type=int)))
+    if jira_id is None:
+        jira_id = click.prompt(
+            'Enter an Airflow JIRA id', value_proc=validate_jira_id)
+    else:
+        jira_id = validate_jira_id(jira_id)
+
     try:
         issue = asf_jira.issue(jira_id)
     except Exception as e:
-        fail("ASF JIRA could not find issue {}\n{}".format(jira_id, e))
-
-    if comment is None:
-        comment = click.prompt(
-            'Please enter a comment to explain why the issue is being closed',
-            default='',
-            show_default=False)
-    if not comment:
-        comment = None
+        raise ValueError(
+            "ASF JIRA could not find issue {}\n{}".format(jira_id, e))
 
     cur_status = issue.fields.status.name
     cur_summary = issue.fields.summary
@@ -338,10 +383,17 @@ def resolve_jira_issue(comment=None, jira_id=None, merge_branches=None):
         cur_assignee = cur_assignee.displayName
 
     if cur_status == "Resolved" or cur_status == "Closed":
-        fail("JIRA issue %s already has status '%s'" % (jira_id, cur_status))
-    click.echo ("=== JIRA %s ===" % jira_id)
+        raise ValueError(
+            "JIRA issue %s already has status '%s'" % (jira_id, cur_status))
+    click.echo ("\n=== JIRA %s ===" % jira_id)
     click.echo ("summary\t\t%s\nassignee\t%s\nstatus\t\t%s\nurl\t\t%s/%s\n" % (
         cur_summary, cur_assignee, cur_status, JIRA_BASE, jira_id))
+    continue_maybe('Proceed with AIRFLOW-{}?'.format(jira_id))
+
+    comment = click.prompt(
+        'Please enter a comment to explain why {jid} is being closed'.format(
+            jira_id),
+        default=comment)
 
     versions = asf_jira.project_versions("AIRFLOW")
     versions = sorted(versions, key=lambda x: x.name, reverse=True)
@@ -428,11 +480,11 @@ def standardize_jira_ref(text):
     components = []
 
     # If the string is compliant, no need to process any further
-    if (re.search(r'^\[AIRFLOW-[0-9]{3,6}\](\[[A-Z0-9_\s,]+\] )+\S+', text)):
+    if (re.search(r'^\[AIRFLOW-[0-9]{1,6}\](\[[A-Z0-9_\s,]+\] )+\S+', text)):
         return text
 
     # Extract JIRA ref(s):
-    pattern = re.compile(r'(AIRFLOW[-\s]*[0-9]{3,6})+', re.IGNORECASE)
+    pattern = re.compile(r'(AIRFLOW[-\s]*[0-9]{1,6})+', re.IGNORECASE)
     for ref in pattern.findall(text):
         # Add brackets, replace spaces with a dash, & convert to uppercase
         jira_refs.append('[' + re.sub(r'\s+', '-', ref.upper()) + ']')
@@ -590,15 +642,20 @@ def main(pr_num, local=False):
     while raw_input("\n%s (y/n): " % pick_prompt).lower() == "y":
         merged_refs = merged_refs + [cherry_pick(pr_num, merge_hash, latest_branch)]
 
-    continue_maybe("Would you like to update an associated JIRA?")
+    continue_maybe("Would you like to update associated JIRA issues?")
     jira_comment = "Issue resolved by pull request %s\n[%s/%s]" % (pr_num, GITHUB_BASE, pr_num)
-    jira_ids = re.findall("AIRFLOW-[0-9]{1,6}", title)
-    if not jira_ids:
-        resolve_jira_issue(
-            jira_id=None, comment=jira_comment, merge_branches=merged_refs)
+
+    jira_ids = re.findall("AIRFLOW-[0-9]{1,6}", title) or [None]
     for jira_id in jira_ids:
         resolve_jira_issue(
-            jira_id=jira_id, comment=jira_comment, merge_branches=merged_refs)
+            jira_id=jira_id,
+            comment=jira_comment,
+            merge_branches=merged_refs)
+
+    if click.confirm('Would you like to resolve another JIRA issue?'):
+        resolve_jira_issues_loop(
+            comment=jira_comment,
+            merged_branches=merged_refs)
 
 
 @click.group()
@@ -645,7 +702,7 @@ def close_jira():
     This command runs only the JIRA part of the PR tool; it doesn't do any
     merging at all.
     """
-    resolve_jira_issue(comment=None, jira_id=None, merge_branches=None)
+    resolve_jira_issues_loop()
 
 
 if __name__ == "__main__":