airflow-pr

#!/usr/bin/env python

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# Utility for creating well-formed pull request merges and pushing them to
# Apache.
#
#   usage: ./airflow-pr    (see config env vars below)
#
# This utility assumes you already have a local Airflow git folder and that you
# have added remotes corresponding to both (i) the github apache Airflow
# mirror and (ii) the apache git repo.

# This tool is based on the Spark merge_spark_pr script:
# https://github.com/apache/spark/blob/master/dev/merge_spark_pr.py

import json
import os
import re
import subprocess
import sys
import textwrap
import urllib.request as urllib

try:
    import click
except ImportError:
    print("Could not find the click library. Run 'sudo pip install click' to install.")
    sys.exit(-1)

try:
    import keyring
except ImportError:
    print("Could not find the keyring library. "
          "Run 'sudo pip install keyring' to install.")
    sys.exit(-1)

# Location of your Airflow git development area
AIRFLOW_GIT_LOCATION = os.environ.get(
    "AIRFLOW_GIT",
    os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

# Remote name which points to the GitHub site
GITHUB_REMOTE_NAME = os.environ.get("GITHUB_REMOTE_NAME", "github")
# OAuth key used for issuing requests against the GitHub API. If this is not
# defined, then requests will be unauthenticated. You should only need to
# configure this if you find yourself regularly exceeding your IP's
# unauthenticated request rate limit. You can create an OAuth key at
# https://github.com/settings/tokens. This tool only requires the "public_repo"
# scope.
GITHUB_OAUTH_KEY = os.environ.get("GITHUB_OAUTH_KEY")

GITHUB_BASE = "https://github.com/apache/airflow/pull"
GITHUB_API_BASE = "https://api.github.com/repos/apache/airflow"
GITHUB_USER = 'asfgit'

JIRA_BASE = "https://issues.apache.org/jira/browse"
JIRA_API_BASE = "https://issues.apache.org/jira"
# Prefix added to temporary branches
BRANCH_PREFIX = "PR_TOOL"

TMP_CREDENTIALS = {}


class PRToolError(Exception):
    pass


def reflow(text, width=80):
    """
    Reformats text so that each line does not exceed `width` characters.

    Preserves any whitespace that follows a newline, such as paragraph breaks
    and indentation.
    """
    new_text = textwrap.dedent(text).strip()
    split = re.split('(\n+\s*)', new_text)
    paragraphs, whitespace = split[::2], split[1::2]
    reformatted = [textwrap.fill(p, width=width) for p in paragraphs]
    result = reformatted + whitespace
    result[::2] = reformatted
    result[1::2] = whitespace
    return ''.join(result)


def get_json(url):
    try:
        request = urllib.Request(url)
        if GITHUB_OAUTH_KEY:
            request.add_header('Authorization', 'token %s' % GITHUB_OAUTH_KEY)

        # decode response for Py3 compatibility
        response = urllib.urlopen(request).read().decode('utf-8')
        return json.loads(response)
    except urllib.HTTPError as e:
        if (
                "X-RateLimit-Remaining" in e.headers and
                e.headers["X-RateLimit-Remaining"] == '0'):
            click.echo(reflow(
                "Exceeded the GitHub API rate limit; set the environment "
                "variable GITHUB_OAUTH_KEY in order to make authenticated "
                "GitHub requests."))
        else:
            click.echo("Unable to fetch URL, exiting: %s" % url)
        sys.exit(-1)


def fail(msg):
    click.echo(msg)
    clean_up()
    sys.exit(-1)


def run_cmd(cmd, echo_cmd=True):
    if isinstance(cmd, list):
        if echo_cmd:
            click.echo('>> Running command: {}'.format(' '.join(cmd)))
        return subprocess.check_output(cmd).decode('utf-8').strip()
    else:
        if echo_cmd:
            click.echo('>> Running command: {}'.format(cmd))
        return subprocess.check_output(cmd.split(" ")).decode('utf-8').strip()


def continue_maybe(prompt):
    if not click.confirm(click.style(prompt, fg='blue', bold=True)):
        fail("Okay, exiting.")


def clean_up():
    if 'original_head' not in globals():
        return

    click.echo('Resetting git to remove any changes')
    run_cmd('git reset --hard')

    click.echo("Restoring head pointer to %s" % original_head)
    run_cmd("git checkout %s" % original_head)

    branches = run_cmd("git branch").replace(" ", "").split("\n")

    for branch in filter(lambda x: x.startswith(BRANCH_PREFIX), branches):
        click.echo("Deleting local branch %s" % branch)
        run_cmd("git branch -D %s" % branch)


# merge the requested PR and return the merge hash
def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, local):

    pr_branch_name = "%s_MERGE_PR_%s" % (BRANCH_PREFIX, pr_num)
    target_branch_name = "%s_MERGE_PR_%s_%s" % (BRANCH_PREFIX, pr_num, target_ref.upper())
    run_cmd("git fetch %s pull/%s/head:%s" % (GITHUB_REMOTE_NAME, pr_num, pr_branch_name))
    run_cmd("git fetch %s %s:%s" % (GITHUB_REMOTE_NAME, target_ref, target_branch_name))
    run_cmd("git checkout %s" % target_branch_name)

    had_conflicts = False
    squash = click.confirm('\n'.join([
        click.style('\nDo you want to squash the PR?\n', bold=True),
        reflow(
            """
            We recommend that you do!

            Squashing will give you an opportunity to edit the new commit
            message, but the squashed commit will still be attributed to the PR
            author. GitHub will show that you merged the squash commit but will
            mark the PR as closed rather than merged (the distinction is purely
            cosmetic).

            If you don't squash, a merge commit will be created in addition to
            the PR commits, but GitHub will properly show the PR as "merged".
            We suggest you do this only if the PR commits are logically
            distinct and should remain separate.
            """),
        click.style('Squash?', fg='blue', bold=True)]),
        default=True)

    if squash:
        merge_cmd = ['git', 'merge', pr_branch_name, '--squash']
    else:
        merge_cmd = ['git', 'merge', pr_branch_name, '--no-ff', '--no-commit']
    try:
        run_cmd(merge_cmd)
    except Exception as e:
        msg = "Error merging: %s\nWould you like to manually fix-up this merge?" % e
        continue_maybe(msg)
        msg = ("Okay, please fix any conflicts and 'git add' conflicting files... " +
               "Finished?")
        continue_maybe(msg)
        had_conflicts = True

    pr_commits = get_json("{}/pulls/{}/commits".format(GITHUB_API_BASE, pr_num))

    # find all JIRA issues mentioned in the text
    all_text = title + body
    if pr_commits:
        all_text += ' '.join(c['commit']['message'] for c in pr_commits)
    all_jira_refs = standardize_jira_ref(all_text, only_jira=True)

    merge_message_flags = []

    if squash:

        # -- create commit message subject
        # if there is only one commit, take the squash commit message from it
        if len(pr_commits) == 1:
            click.echo(click.style('\n' + reflow(
                """
                This squash contains only one commit, so we will use its
                commit message for the squash commit message. We will
                automatically add references to every JIRA issue
                mentioned in the PR. You will have an opportunity to edit
                it later.
                """), bold=True))
            commit_message = pr_commits[0]['commit']['message']
            merge_message_flags.extend(["-m", commit_message])
        # if there is are multiple commits, take the squash commit message from
        # the PR title
        else:
            click.echo(click.style('\n' + reflow(
                """
                This squash contains more than one commit, so we will use
                the PR title as the squash commit subject. We will
                automatically add references to every JIRA issue mentioned in
                the PR. You will have an opportunity to edit it later.
                """), bold=True))
            merge_message_flags.extend(["-m", title])

        # add all JIRA refs to the commit subject and then standardize it
        # to get a clean reference to every JIRA issue mentioned in the PR
        first_commit_msg = merge_message_flags[-1].split('\n')
        first_commit_msg[0] = standardize_jira_ref(
            first_commit_msg[0] + all_jira_refs)
        if not re.findall("AIRFLOW-[0-9]{1,6}", first_commit_msg[0]):
            continue_maybe(click.style('\n' + reflow(
                """
                This PR doesn't reference any JIRA issues!!

                Are you sure you want to continue?
                """), fg='red', bold=True))
        merge_message_flags[-1] = '\n'.join(first_commit_msg)

        # -- Note conflicts
        if had_conflicts:
            committer_name = run_cmd(
                "git config --get user.name", echo_cmd=False)
            committer_email = run_cmd(
                "git config --get user.email", echo_cmd=False)
            message = (
                'This patch had conflicts when merged, resolved by '
                'Committer: %s <%s>' % (committer_name, committer_email))
            merge_message_flags.extend(["-m", message])

        # -- Add PR body to commit message
        msg = click.style(
            '\nWould you like to include the PR body in the squash '
            'commit message?',
            fg='blue', bold=True)
        if body and click.confirm(msg, default=False, prompt_suffix=''):
            # We remove @ symbols from the body to avoid triggering e-mails
            # to people every time someone creates a public fork of Airflow.
            merge_message_flags += ["-m", body.replace("@", "")]

        # -- add individual commit messages to squash commit
        if len(pr_commits) > 1:
            m = click.style(
                'Would you like to include the individual commit messages '
                'in the squash commit message?',
                fg='blue', bold=True)
            if pr_commits and click.confirm(m, default=True, prompt_suffix=''):
                for commit in pr_commits:
                    merge_message_flags.extend(
                        ['-m', commit['commit']['message']])

        # The string "Closes #%s" string is required for GitHub to correctly
        # close the PR. GitHub will mark the PR as closed, not merged
        close_msg = "closes #{}".format(pr_num)
        merge_message_flags.extend(["-m", "{} from {}".format(
            close_msg.capitalize(), pr_repo_desc)])

        # -- set authors and add authors to commit message
        commit_authors = run_cmd(
            ['git', 'log', 'HEAD..{}'.format(pr_branch_name),
             '--pretty=format:%an <%ae>'], echo_cmd=False).split("\n")
        distinct_authors = sorted(
            set(commit_authors),
            key=lambda x: commit_authors.count(x),
            reverse=True)
        primary_author = click.prompt(
            click.style(
                'Enter the primary author in the format of \"name <email>\"',
                fg='blue', bold=True),
            default=distinct_authors[0])
        if primary_author == "":
            primary_author = distinct_authors[0]

        merge_message_flags.append('--author="{}"'.format(primary_author))

    else:
        # This will mark the PR as merged
        merge_message_flags.extend([
            '-m',
            'Merge pull request #{} from {}'.format(pr_num, pr_repo_desc)])

    # reflow commit message
    seen_first_line = False
    for i in range(1, len(merge_message_flags)):
        if merge_message_flags[i - 1] == '-m':
            # let the first line be as long as the user wants
            if not seen_first_line:
                if '\n\n' in merge_message_flags[i]:
                    title, body = merge_message_flags[i].split('\n\n', 1)
                    body = reflow(body, 50)
                    merge_message_flags[i] = title + '\n\n' + body
                seen_first_line = True
            else:
                merge_message_flags[i] = reflow(merge_message_flags[i], 50)

    run_cmd(['git', 'commit'] + merge_message_flags, echo_cmd=False)

    if squash:
        # -- ask user to edit commit message
        click.echo(click.style('\n=== Current Squash Commit ===', bold=True))
        click.echo(run_cmd('git log -1 --pretty=%B', echo_cmd=False))
        click.echo(click.style('=== End of Squash Commit ===\n', bold=True))
        msg = reflow(
            """
            If you would like to edit the commit message, open a new
            terminal and run:

                git commit --amend

            When you have finished, return here and press any key to
            continue.
            """)
        click.pause(click.style(msg, fg='blue', bold=True))

        # The user might have removed "Closes #XXXX" from the commit message
        # so we add it back to make sure GitHub closes the PR.
        commit_msg = run_cmd('git log -1 --pretty=%B', echo_cmd=False)
        if close_msg not in commit_msg.lower():
            click.echo(reflow("""
                Your commit message does not contain the phrase "{}".
                Without it, GitHub can\'t link this commit to the PR. We
                will automatically add it to the end of your commit
                message.""".format(close_msg)))
            commit_flags = []
            commit_flags.append('--author="{}"'.format(primary_author))
            commit_flags.extend(['-m', commit_msg])
            commit_flags.extend(
                ["-m", "{} from {}".format(
                    close_msg.capitalize(), pr_repo_desc)])
            run_cmd('git reset --soft HEAD~1', echo_cmd=False)
            run_cmd(['git', 'commit'] + commit_flags, echo_cmd=False)

    if local:
        msg = '\n' + reflow("""
            The PR has been merged locally in branch {}.
            You may leave this program running while you work on it. When
            you are finished, press any key to delete the PR branch and
            restore your original environment.
            """.format(target_branch_name))

        click.pause(click.style(msg, fg='blue', bold=True))

        clean_up()
        return
    else:
        continue_maybe(
            '\n\nThe local merge is complete ({}).\n'.format(
                target_branch_name) +
            click.style(
                'Push to Gitbox ({})?'.format(GITHUB_REMOTE_NAME), 'red'))

    try:
        run_cmd('git push %s %s:%s' % (
            GITHUB_REMOTE_NAME, target_branch_name, target_ref))
    except Exception as e:
        clean_up()
        fail("Exception while pushing: %s" % e)

    merge_hash = run_cmd("git rev-parse %s" % target_branch_name)[:8]
    clean_up()
    click.echo("Pull request #%s merged!" % pr_num)
    click.echo("Merge hash: %s" % merge_hash)
    return merge_hash


def cherry_pick(pr_num, merge_hash, default_branch):
    pick_ref = click.prompt(click.style(
        "Enter a branch name (or press enter to use %s): " % default_branch,
        fg='blue', bold=True))
    if pick_ref == "":
        pick_ref = default_branch

    pick_branch_name = "%s_PICK_PR_%s_%s" % (
        BRANCH_PREFIX, pr_num, pick_ref.upper())

    run_cmd("git fetch %s %s:%s" % (
        GITHUB_REMOTE_NAME, pick_ref, pick_branch_name))
    run_cmd("git checkout %s" % pick_branch_name)

    try:
        run_cmd("git cherry-pick -sx %s" % merge_hash)
    except Exception as e:
        msg = (
            "Error cherry-picking: {}\n"
            "Would you like to manually fix-up this merge?".format(e))
        continue_maybe(msg)
        msg = (
            "Okay, please fix any conflicts and finish the cherry-pick. "
            "Finished?")
        continue_maybe(msg)

    continue_maybe("Pick complete (local ref %s). Push to %s?" % (
        pick_branch_name, GITHUB_REMOTE_NAME))

    try:
        run_cmd(
            'git push %s %s:%s' % (
                GITHUB_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]
    clean_up()

    click.echo("Pull request #%s picked into %s!" % (pr_num, pick_ref))
    click.echo("Pick hash: %s" % pick_hash)
    return pick_ref


def fix_version_from_branch(branch, versions):
    # Note: Assumes this is a sorted (newest->oldest) list of un-released
    # versions
    if branch == "master":
        return versions[0]
    else:
        # TODO adopt a release scheme with branches. Spark uses branch-XX.
        branch_ver = branch.replace("branch-", "")
        versions = list(filter(
            lambda x: x.name.startswith(branch_ver), versions))
        if versions:
            return versions[-1]


def register(username, password):
    """ Use this function to register a JIRA account in your OS' keyring """
    keyring.set_password('airflow-pr', 'username', username)
    keyring.set_password('airflow-pr', 'password', password)


def validate_jira_id(jira_id):
    if not jira_id:
        return
    elif isinstance(jira_id, int):
        return 'AIRFLOW-{}'.format(abs(jira_id))

    # 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 PRToolError as e:
            click.echo("PR Tool Error: {}".format(e))
            sys.exit(-1)
        except Exception as e:
            click.echo("ERROR: {}".format(e))

        if not click.confirm(click.style(
                'Would you like to resolve another JIRA issue?',
                fg='blue', bold=True)):
            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.
    """
    try:
        import jira.client
    except ImportError:
        raise PRToolError(
            "Could not find jira-python library; exiting. Run "
            "'sudo pip install jira' to install.")

    if merge_branches is None:
        merge_branches = []

    # ASF JIRA username
    JIRA_USERNAME = os.environ.get("JIRA_USERNAME", '')
    if not JIRA_USERNAME:
        JIRA_USERNAME = TMP_CREDENTIALS.get('JIRA_USERNAME', '')
    if not JIRA_USERNAME:
        JIRA_USERNAME = keyring.get_password("airflow-pr", "username")
        if JIRA_USERNAME:
            click.echo("Obtained jira username from keyring. To reset remove it there")

    # ASF JIRA password
    JIRA_PASSWORD = os.environ.get("JIRA_PASSWORD", '')
    if not JIRA_PASSWORD:
        JIRA_PASSWORD = TMP_CREDENTIALS.get('JIRA_PASSWORD', '')

    if not JIRA_USERNAME:
        JIRA_USERNAME = click.prompt(
            click.style('Username for Airflow JIRA', fg='blue', bold=True),
            type=str)
        click.echo(
            'Set a JIRA_USERNAME env var to avoid this prompt in the future.')
        TMP_CREDENTIALS['JIRA_USERNAME'] = JIRA_USERNAME
    if JIRA_USERNAME and not JIRA_PASSWORD:
        JIRA_PASSWORD = keyring.get_password("airflow-pr", 'password')
        if JIRA_PASSWORD:
            click.echo("Obtained password from keyring. To reset remove it there.")
    if not JIRA_PASSWORD:
        JIRA_PASSWORD = click.prompt(
            click.style('Password for Airflow JIRA', fg='blue', bold=True),
            type=str,
            hide_input=True)
        if JIRA_USERNAME and JIRA_PASSWORD:
            if click.confirm(click.style("Would you like to store your password "
                                         "in your keyring?", fg='blue', bold=True)):
                register(JIRA_USERNAME, JIRA_PASSWORD)
        TMP_CREDENTIALS['JIRA_PASSWORD'] = JIRA_PASSWORD

    try:
        asf_jira = jira.client.JIRA(
            {'server': JIRA_API_BASE},
            basic_auth=(JIRA_USERNAME, JIRA_PASSWORD))
    except:
        raise ValueError('Could not log in to JIRA!')

    if jira_id is None:
        jira_id = click.prompt(click.style(
            'Enter an Airflow JIRA id', fg='blue', bold=True),
            value_proc=validate_jira_id)
    else:
        jira_id = validate_jira_id(jira_id)

    try:
        issue = asf_jira.issue(jira_id)
    except Exception as e:
        raise ValueError(
            "ASF JIRA could not find issue {}\n{}".format(jira_id, e))

    cur_status = issue.fields.status.name
    cur_summary = issue.fields.summary
    cur_assignee = issue.fields.assignee
    if cur_assignee is None:
        cur_assignee = "NOT ASSIGNED!!!"
    else:
        cur_assignee = cur_assignee.displayName

    # check if issue was already closed
    if cur_status == "Resolved" or cur_status == "Closed":
        click.echo("JIRA issue {} already has status '{}'".format(
            jira_id, cur_status))
        return

    click.echo(click.style("\n === JIRA %s ===" % jira_id, bold=True))
    click.echo(
        "summary:\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))
    if not click.confirm(click.style(
            'Proceed with {}?'.format(jira_id), fg='blue', bold=True)):
        return

    if comment is None:
        comment = click.prompt(
            click.style(
                'Please enter a comment to explain why this issue '
                'is being closed',
                fg='blue', bold=True),
            default='',
            show_default=False)

    versions = asf_jira.project_versions("AIRFLOW")
    versions = sorted(versions, key=lambda x: x.name, reverse=True)
    versions = filter(lambda x: x.raw['released'] is False, versions)
    # Consider only x.y.z versions
    versions = list(filter(
        lambda x: re.match('\d+\.\d+\.\d+', x.name), versions))

    if versions:
        default_fix_versions = map(
            lambda x: fix_version_from_branch(x, versions), merge_branches)
        default_fix_versions = [v.name for v in default_fix_versions if v]
    else:
        default_fix_versions = []

    # TODO Airflow versions vary from two to four decimal places (2.0, 1.7.1.3)
    # The following logic can be reintroduced if/when a standard emerges.

    # for v in default_fix_versions:

    #     Handles the case where we have forked a release branch but not yet
    #     made the release. In this case, if the PR is committed to the master
    #     branch and the release branch, we only consider the release branch to
    #     be the fix version. E.g. it is not valid to have both 1.1.0 and 1.0.0
    #     as fix versions.

    #     (major, minor, patch) = v.split(".")
    #     if patch == "0":
    #         previous = "%s.%s.%s" % (major, int(minor) - 1, 0)
    #         if previous in default_fix_versions:
    #             default_fix_versions = list(filter(
    #                 lambda x: x != v, default_fix_versions))
    default_fix_versions = ",".join(default_fix_versions)

    fix_versions = click.prompt(
        click.style(
            "Enter comma-separated fix version(s)", fg='blue', bold=True),
        default=default_fix_versions)
    if fix_versions == "":
        fix_versions = default_fix_versions
    fix_versions = fix_versions.replace(" ", "").split(",")
    if fix_versions == ['']:
        fix_versions = None

    def get_version_json(version_str):
        version_list = list(filter(lambda v: v.name == version_str, versions))
        if version_list:
            return version_list[0].raw
        else:
            return ''

    if fix_versions and fix_versions != ['']:
        jira_fix_versions = list(map(get_version_json, fix_versions))
    else:
        jira_fix_versions = None

    action = list(filter(
        lambda a: a['name'] == 'Resolve Issue',
        asf_jira.transitions(jira_id)))[0]
    resolution = list(filter(
        lambda r: r.raw['name'] == "Fixed",
        asf_jira.resolutions()))[0]
    asf_jira.transition_issue(
        jira_id,
        action["id"],
        fixVersions=jira_fix_versions,
        comment=comment or None,
        resolution={'id': resolution.raw['id']})

    click.echo("Successfully resolved {id}{fv}!".format(
        id=jira_id,
        fv=' with fix versions={}'.format(fix_versions) if fix_versions else ''
    ))


def standardize_jira_ref(text, only_jira=False):
    """
    Standardize the [AIRFLOW-XXXXX] [MODULE] prefix
    Converts "[AIRFLOW-XXX][mllib] Issue", "[MLLib] AIRFLOW-XXX. Issue" or
    "AIRFLOW XXX [MLLIB]: Issue" to "[AIRFLOW-XXX][MLLIB] Issue"

    >>> standardize_jira_ref("[AIRFLOW-5821] [SQL] ParquetRelation2 CTAS should check if delete is successful")
    '[AIRFLOW-5821][SQL] ParquetRelation2 CTAS should check if delete is successful'
    >>> standardize_jira_ref("[AIRFLOW-4123][Project Infra][WIP]: Show new dependencies added in pull requests")
    '[AIRFLOW-4123][PROJECT INFRA][WIP] Show new dependencies added in pull requests'
    >>> standardize_jira_ref("[MLlib] Airflow  5954: Top by key")
    '[AIRFLOW-5954][MLLIB] Top by key'
    >>> standardize_jira_ref("[AIRFLOW-979] a LRU scheduler for load balancing in TaskSchedulerImpl")
    '[AIRFLOW-979] a LRU scheduler for load balancing in TaskSchedulerImpl'
    >>> standardize_jira_ref("AIRFLOW-1094 Support MiMa for reporting binary compatibility across versions.")
    '[AIRFLOW-1094] Support MiMa for reporting binary compatibility across versions.'
    >>> standardize_jira_ref("[WIP]  [AIRFLOW-1146] Vagrant support for Spark")
    '[AIRFLOW-1146][WIP] Vagrant support for Spark'
    >>> standardize_jira_ref("AIRFLOW-1032. If Yarn app fails before registering, app master stays aroun...")
    '[AIRFLOW-1032] If Yarn app fails before registering, app master stays aroun...'
    >>> standardize_jira_ref("[AIRFLOW-6250][AIRFLOW-6146][AIRFLOW-5911][SQL] Types are now reserved words in DDL parser.")
    '[AIRFLOW-6250][AIRFLOW-6146][AIRFLOW-5911][SQL] Types are now reserved words in DDL parser.'
    >>> standardize_jira_ref("Additional information for users building from source code")
    'Additional information for users building from source code'
    >>> standardize_jira_ref('AIRFLOW 35 AIRFLOW--36 AIRFLOW  37 test', only_jira=True)
    '[AIRFLOW-35][AIRFLOW-36][AIRFLOW-37]'
    """  # noqa
    jira_refs = []
    components = []

    pattern = re.compile(r'([\[]*AIRFLOW[-\s]*[0-9]{1,6}[\]]*)', re.IGNORECASE)
    for ref in pattern.findall(text):

        orig_ref = ref
        if not re.findall("[AIRFLOW-[0-9]{1,6}]", ref):
            # convert to uppercase
            ref = ref.upper()
            # replace 0+ spaces with a dash
            ref = re.sub(r'(AIRFLOW)[-\s]*([0-9]{1,6})', r'\1-\2', ref.upper())

            # and add brackets if needed
            if not ref.startswith('['):
                ref = '[' + ref
            if not ref.endswith(']'):
                ref = ref + ']'

        jira_refs.append(ref)
        text = text.replace(orig_ref, '')

    # Extract Airflow component(s):
    # Look for alphanumeric chars, spaces, dashes, periods, and/or commas
    pattern = re.compile(r'(\[[\w\s,-\.]+\])', re.IGNORECASE)
    for component in pattern.findall(text):
        components.append(component.upper())
        text = text.replace(component, '')

    # Cleanup any remaining symbols:
    pattern = re.compile(r'^\W+(.*)', re.IGNORECASE)
    if pattern.search(text) is not None:
        text = pattern.search(text).groups()[0]

    def unique(seq):
        new_seq = []
        for s in seq:
            if s not in new_seq:
                new_seq.append(s)
        return new_seq

    # Assemble full text (JIRA ref(s), module(s), remaining text)
    clean_text = ''.join(unique(jira_refs)).strip()
    if not only_jira:
        clean_text += (
            ''.join(unique(components)).strip() + " " + text.strip())

    # Replace multiple spaces with a single space, e.g. if no jira refs
    # and/or components were included
    clean_text = re.sub(r'\s+', ' ', clean_text.strip())

    return clean_text


def get_current_ref():
    ref = run_cmd("git rev-parse --abbrev-ref HEAD")
    if ref == 'HEAD':
        # The current ref is a detached HEAD, so grab its SHA.
        return run_cmd("git rev-parse HEAD")
    else:
        return ref


def main(pr_num, local=False):
    """
    Utility for creating well-formed pull request merges and pushing them
    to Apache.

    This tool assumes you already have a local Airflow git folder and that you
    have added remotes corresponding to both (i) the github apache Airflow
    mirror and (ii) the apache git repo.

    To configure the tool, set the following env vars:
    - AIRFLOW_GIT
        The location of your Airflow git development area (defaults to the
        current working directory)

    - GITHUB_REMOTE_NAME
        GitHub remote name (defaults to "github")

    - JIRA_USERNAME
        ASF JIRA username for automatically closing JIRA issues. Users will be
        prompted if it is not set.

    - JIRA_PASSWORD
        ASF JIRA password for automatically closing JIRA issues. Users will be
        prompted if it is not set.

    - GITHUB_OAUTH_KEY
        Only required if you are exceeding the rate limit for a single IP
        address.
    """
    global original_head

    os.chdir(AIRFLOW_GIT_LOCATION)
    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])
    # Assumes branch names can be sorted lexicographically
    latest_branch = sorted(branch_names, reverse=True)
    if latest_branch:
        latest_branch = latest_branch[0]
    else:
        latest_branch = ''

    if not pr_num:
        pr_num = click.prompt(
            click.style(
                "Please enter the number of the pull request you'd "
                "like to work with",
                fg='blue', bold=True),
            type=int)
    else:
        click.echo('Working with pull request {}'.format(pr_num))

    pr = get_json("{}/pulls/{}".format(GITHUB_API_BASE, pr_num))

    url = pr["url"]

    # Decide whether to use the modified title or not
    modified_title = standardize_jira_ref(pr["title"])
    if modified_title != pr["title"]:
        click.echo("I've re-written the title to match the standard format:")
        click.echo("Original: %s" % pr["title"])
        click.echo("Modified: %s" % modified_title)
        result = click.confirm(click.style(
            "Would you like to use the modified title?", fg='blue', bold=True))
        if result:
            title = modified_title
            click.echo("Using modified title:")
        else:
            title = pr["title"]
            click.echo("Using original title:")
        click.echo(title)
    else:
        title = pr["title"]

    body = pr["body"]
    target_ref = pr["base"]["ref"]
    user_login = pr["user"]["login"]
    base_ref = pr["head"]["ref"]
    pr_repo_desc = "%s/%s" % (user_login, base_ref)

    if not bool(pr["mergeable"]):
        msg = ('Pull request {} is not mergeable in its current form.\n'
               'Continue anyway? (experts only!)'.format(pr_num))
        continue_maybe(msg)

    click.echo(click.style("\n=== Pull Request #%s ===" % pr_num, bold=True))
    click.echo("title:\t%s\nsource:\t%s\ntarget:\t%s\nurl:\t%s\n" % (
        title, pr_repo_desc, target_ref, url))
    continue_maybe("Proceed with pull request #{}?".format(pr_num))

    merged_refs = [target_ref]

    merge_hash = merge_pr(pr_num, target_ref, title, body, pr_repo_desc, local)

    if local:
        return

    msg = "Would you like to pick {} into another branch?".format(merge_hash)
    while click.confirm(click.style(msg, fg='blue', bold=True)):
        merged_refs = merged_refs + [
            cherry_pick(pr_num, merge_hash, latest_branch)]

    jira_ids = set(re.findall("AIRFLOW-[0-9]{1,6}", title + body) or [None])

    msg = reflow(
        """
        We found {n} JIRA issue{s} referenced in the PR. Would you
        like to update{it}{them} any other JIRA issues?
        """.format(
            n=len(jira_ids) if jira_ids else 'no',
            s='s' if len(jira_ids) != 1 else '',
            it=' it or' if len(jira_ids) == 1 else '',
            them=' them or' if len(jira_ids) > 1 else ''))
    if not click.confirm(click.style(msg, fg='blue', bold=True), default=True):
        fail("Okay, exiting.")

    jira_comment = "Issue resolved by pull request #{}\n[{}/{}]".format(
        pr_num, GITHUB_BASE, pr_num)

    for jira_id in set(jira_ids):
        resolve_jira_issue(
            jira_id=jira_id,
            comment=jira_comment,
            merge_branches=merged_refs)

    if not jira_ids or click.confirm(click.style(
            'Would you like to resolve another JIRA issue?',
            fg='blue', bold=True)):
        resolve_jira_issues_loop(
            comment=jira_comment,
            merge_branches=merged_refs)


@click.group()
def cli():
    r"""
    This tool should be used by Airflow committers to test PRs, merge them
    into the master branch, and close related JIRA issues.

    Before you begin, make sure you have created the 'github' git remote. You
    can use the "setup_git_remotes" command to do this automatically. If you do
    not want to use these remote names, you can tell the PR tool by setting the
    appropriate environment variable. For more information, run:

        airflow-pr merge --help
    """
    os.chdir(AIRFLOW_GIT_LOCATION)
    status = run_cmd('git status --porcelain', echo_cmd=False)
    if status:
        msg = (
            'You have uncommitted changes in this branch. Running this tool\n'
            'will delete them permanently. Continue?')
        if click.confirm(click.style(msg, fg='red', bold=True)):
            run_cmd('git reset --hard', echo_cmd=False)
        else:
            sys.exit(-1)


@cli.command(short_help='Merge a GitHub PR into Airflow master')
@click.argument('pr_num', default=0)
def merge(pr_num):
    """
    Utility for creating well-formed pull request merges and pushing them
    to Gitbox (a.k.a. GitHub), as well as closing JIRA issues.

    This tool assumes you already have a local Airflow git folder and that you
    have added remote corresponding to the github apache Airflow repo.

    To configure the tool, set the following env vars:
    - AIRFLOW_GIT
        The location of your Airflow git development area (defaults to the
        current working directory)

    - GITHUB_REMOTE_NAME
        GitHub remote name (defaults to "github")

    - JIRA_USERNAME
        ASF JIRA username for automatically closing JIRA issues. Users will be
        prompted if it is not set.

    - JIRA_PASSWORD
        ASF JIRA password for automatically closing JIRA issues. Users will be
        prompted if it is not set.

    - GITHUB_OAUTH_KEY
        Only required if you are exceeding the rate limit for a single IP
        address.
    """
    main(pr_num, local=False)


@cli.command(short_help='Clone a GitHub PR locally for testing (no push)')
@click.argument('pr_num', default=0)
def work_local(pr_num):
    """
    Clones a PR locally for testing, imitating a full merge workflow, but does
    not push the changes to Airflow master. Instead, the program will pause
    once the local merge is complete, allowing the user to explore any changes.
    Once finished, the program will delete the merge and restore the original
    environment.
    """
    main(pr_num, local=True)


@cli.command(short_help='Close a JIRA issue (without merging a PR)')
def close_jira():
    """
    This command runs only the JIRA part of the PR tool; it doesn't do any
    merging at all.
    """
    resolve_jira_issues_loop()


@cli.command(short_help='Set up default git remotes')
def setup_git_remotes():
    click.echo(reflow("""
        This command will create git remotes to mirror the following
        structure. If you do not want to use these names, you must set the
        GITHUB_REMOTE_NAME environment variable:

          git remote -v
          github https://github.com/apache/airflow.git (fetch)
          github https://github.com/apache/airflow.git (push)

        If these remotes already exist, the tool will display an error.
        """))
    continue_maybe('Do you want to continue?')

    error = False
    try:
        run_cmd('git remote add github git@github.com:apache/airflow.git')
    except:
        click.echo(click.style(reflow(
            '>>ERROR: Could not create github remote. If it already exists, '
            'run `git remote remove github` to delete it.', fg='red')))
        error = True
    if not error:
        click.echo('Done setting up git remotes. Run git remote -v to see them.')


if __name__ == "__main__":
    import doctest
    (failure_count, test_count) = doctest.testmod()
    if failure_count:
        exit(-1)
    try:
        cli()
    except:
        clean_up()
        raise