ClickHouse/tests/ci/cherry_pick_utils/query.py

# -*- coding: utf-8 -*-

import requests


class Query:
    '''
    Implements queries to the Github API using GraphQL
    '''

    _PULL_REQUEST = '''
        author {{
            ... on User {{
                id
                login
            }}
        }}

        baseRepository {{
            nameWithOwner
        }}

        mergeCommit {{
            oid
            parents(first: {min_page_size}) {{
                totalCount
                nodes {{
                    oid
                }}
            }}
        }}

        mergedBy {{
            ... on User {{
                id
                login
            }}
        }}

        baseRefName
        closed
        headRefName
        id
        mergeable
        merged
        number
        title
        url
    '''

    def __init__(self, token, owner, name, team, max_page_size=100, min_page_size=10):
        self._PULL_REQUEST = Query._PULL_REQUEST.format(min_page_size=min_page_size)

        self._token = token
        self._owner = owner
        self._name = name
        self._team = team

        self._max_page_size = max_page_size
        self._min_page_size = min_page_size

        self.api_costs = {}

        repo = self.get_repository()
        self._id = repo['id']
        self.ssh_url = repo['sshUrl']
        self.default_branch = repo['defaultBranchRef']['name']

        self.members = set(self.get_members())

    def get_repository(self):
        _QUERY = '''
            repository(owner: "{owner}" name: "{name}") {{
                defaultBranchRef {{
                    name
                }}
                id
                sshUrl
            }}
        '''

        query = _QUERY.format(owner=self._owner, name=self._name)
        return self._run(query)['repository']

    def get_members(self):
        '''Get all team members for organization

        Returns:
            members: a map of members' logins to ids
        '''

        _QUERY = '''
            organization(login: "{organization}") {{
                team(slug: "{team}") {{
                    members(first: {max_page_size} {next}) {{
                        pageInfo {{
                            hasNextPage
                            endCursor
                        }}
                        nodes {{
                            id
                            login
                        }}
                    }}
                }}
            }}
        '''

        members = {}
        not_end = True
        query = _QUERY.format(organization=self._owner, team=self._team,
                              max_page_size=self._max_page_size,
                              next='')

        while not_end:
            result = self._run(query)['organization']['team']
            if result is None:
                break
            result = result['members']
            not_end = result['pageInfo']['hasNextPage']
            query = _QUERY.format(organization=self._owner, team=self._team,
                                  max_page_size=self._max_page_size,
                                  next='after: "{}"'.format(result["pageInfo"]["endCursor"]))

            members += dict([(node['login'], node['id']) for node in result['nodes']])

        return members

    def get_pull_request(self, number):
        _QUERY = '''
            repository(owner: "{owner}" name: "{name}") {{
                pullRequest(number: {number}) {{
                    {pull_request_data}
                }}
            }}
        '''

        query = _QUERY.format(owner=self._owner, name=self._name, number=number,
                              pull_request_data=self._PULL_REQUEST, min_page_size=self._min_page_size)
        return self._run(query)['repository']['pullRequest']

    def find_pull_request(self, base, head):
        _QUERY = '''
            repository(owner: "{owner}" name: "{name}") {{
                pullRequests(first: {min_page_size} baseRefName: "{base}" headRefName: "{head}") {{
                    nodes {{
                        {pull_request_data}
                    }}
                    totalCount
                }}
            }}
        '''

        query = _QUERY.format(owner=self._owner, name=self._name, base=base, head=head,
                              pull_request_data=self._PULL_REQUEST, min_page_size=self._min_page_size)
        result = self._run(query)['repository']['pullRequests']
        if result['totalCount'] > 0:
            return result['nodes'][0]
        else:
            return {}

    def find_pull_requests(self, label_name):
        '''
        Get all pull-requests filtered by label name
        '''
        _QUERY = '''
            repository(owner: "{owner}" name: "{name}") {{
                pullRequests(first: {min_page_size} labels: "{label_name}" states: OPEN) {{
                    nodes {{
                        {pull_request_data}
                    }}
                }}
            }}
        '''

        query = _QUERY.format(owner=self._owner, name=self._name, label_name=label_name,
                              pull_request_data=self._PULL_REQUEST, min_page_size=self._min_page_size)
        return self._run(query)['repository']['pullRequests']['nodes']

    def get_pull_requests(self, before_commit):
        '''
        Get all merged pull-requests from the HEAD of default branch to the last commit (excluding)
        '''

        _QUERY = '''
            repository(owner: "{owner}" name: "{name}") {{
                defaultBranchRef {{
                    target {{
                        ... on Commit {{
                            history(first: {max_page_size} {next}) {{
                                pageInfo {{
                                    hasNextPage
                                    endCursor
                                }}
                                nodes {{
                                    oid
                                    associatedPullRequests(first: {min_page_size}) {{
                                        totalCount
                                        nodes {{
                                            ... on PullRequest {{
                                                {pull_request_data}

                                                labels(first: {min_page_size}) {{
                                                    totalCount
                                                    pageInfo {{
                                                        hasNextPage
                                                        endCursor
                                                    }}
                                                    nodes {{
                                                        name
                                                        color
                                                    }}
                                                }}
                                            }}
                                        }}
                                    }}
                                }}
                            }}
                        }}
                    }}
                }}
            }}
        '''

        pull_requests = []
        not_end = True
        query = _QUERY.format(owner=self._owner, name=self._name,
                              max_page_size=self._max_page_size,
                              min_page_size=self._min_page_size,
                              pull_request_data=self._PULL_REQUEST,
                              next='')

        while not_end:
            result = self._run(query)['repository']['defaultBranchRef']['target']['history']
            not_end = result['pageInfo']['hasNextPage']
            query = _QUERY.format(owner=self._owner, name=self._name,
                                  max_page_size=self._max_page_size,
                                  min_page_size=self._min_page_size,
                                  pull_request_data=self._PULL_REQUEST,
                                  next='after: "{}"'.format(result["pageInfo"]["endCursor"]))

            for commit in result['nodes']:
                # FIXME: maybe include `before_commit`?
                if str(commit['oid']) == str(before_commit):
                    not_end = False
                    break

                # TODO: fetch all pull-requests that were merged in a single commit.
                assert commit['associatedPullRequests']['totalCount'] <= self._min_page_size

                for pull_request in commit['associatedPullRequests']['nodes']:
                    if(pull_request['baseRepository']['nameWithOwner'] == '{}/{}'.format(self._owner, self._name) and
                       pull_request['baseRefName'] == self.default_branch and
                       pull_request['mergeCommit']['oid'] == commit['oid']):
                        pull_requests.append(pull_request)

        return pull_requests

    def create_pull_request(self, source, target, title, description="", draft=False, can_modify=True):
        _QUERY = '''
            createPullRequest(input: {{
                baseRefName: "{target}",
                headRefName: "{source}",
                repositoryId: "{id}",
                title: "{title}",
                body: "{body}",
                draft: {draft},
                maintainerCanModify: {modify}
            }}) {{
                pullRequest {{
                    {pull_request_data}
                }}
            }}
        '''

        query = _QUERY.format(target=target, source=source, id=self._id, title=title, body=description,
                              draft="true" if draft else "false", modify="true" if can_modify else "false",
                              pull_request_data=self._PULL_REQUEST)
        return self._run(query, is_mutation=True)['createPullRequest']['pullRequest']

    def merge_pull_request(self, id):
        _QUERY = '''
            mergePullRequest(input: {{
                pullRequestId: "{id}"
            }}) {{
                pullRequest {{
                    {pull_request_data}
                }}
            }}
        '''

        query = _QUERY.format(id=id, pull_request_data=self._PULL_REQUEST)
        return self._run(query, is_mutation=True)['mergePullRequest']['pullRequest']

    # FIXME: figure out how to add more assignees at once
    def add_assignee(self, pr, assignee):
        _QUERY = '''
            addAssigneesToAssignable(input: {{
                assignableId: "{id1}",
                assigneeIds: "{id2}"
            }}) {{
                clientMutationId
            }}
        '''

        query = _QUERY.format(id1=pr['id'], id2=assignee['id'])
        self._run(query, is_mutation=True)

    def set_label(self, pull_request, label_name):
        '''
        Set label by name to the pull request

        Args:
            pull_request: JSON object returned by `get_pull_requests()`
            label_name (string): label name
        '''

        _GET_LABEL = '''
            repository(owner: "{owner}" name: "{name}") {{
                labels(first: {max_page_size} {next} query: "{label_name}") {{
                    pageInfo {{
                        hasNextPage
                        endCursor
                    }}
                    nodes {{
                        id
                        name
                        color
                    }}
                }}
            }}
        '''

        _SET_LABEL = '''
            addLabelsToLabelable(input: {{
                labelableId: "{pr_id}",
                labelIds: "{label_id}"
            }}) {{
                clientMutationId
            }}
        '''

        labels = []
        not_end = True
        query = _GET_LABEL.format(owner=self._owner, name=self._name, label_name=label_name,
                                  max_page_size=self._max_page_size,
                                  next='')

        while not_end:
            result = self._run(query)['repository']['labels']
            not_end = result['pageInfo']['hasNextPage']
            query = _GET_LABEL.format(owner=self._owner, name=self._name, label_name=label_name,
                                      max_page_size=self._max_page_size,
                                      next='after: "{}"'.format(result["pageInfo"]["endCursor"]))

            labels += [label for label in result['nodes']]

        if not labels:
            return

        query = _SET_LABEL.format(pr_id=pull_request['id'], label_id=labels[0]['id'])
        self._run(query, is_mutation=True)

    def _run(self, query, is_mutation=False):
        from requests.adapters import HTTPAdapter
        from urllib3.util.retry import Retry

        def requests_retry_session(
            retries=3,
            backoff_factor=0.3,
            status_forcelist=(500, 502, 504),
            session=None,
        ):
            session = session or requests.Session()
            retry = Retry(
                total=retries,
                read=retries,
                connect=retries,
                backoff_factor=backoff_factor,
                status_forcelist=status_forcelist,
            )
            adapter = HTTPAdapter(max_retries=retry)
            session.mount('http://', adapter)
            session.mount('https://', adapter)
            return session

        headers = {'Authorization': 'bearer {}'.format(self._token)}
        if is_mutation:
            query = '''
            mutation {{
                {query}
            }}
            '''.format(query=query)
        else:
            query = '''
            query {{
                {query}
                rateLimit {{
                    cost
                    remaining
                }}
            }}
            '''.format(query=query)

        while True:
            request = requests_retry_session().post('https://api.github.com/graphql', json={'query': query}, headers=headers)
            if request.status_code == 200:
                result = request.json()
                if 'errors' in result:
                    raise Exception('Errors occurred: {}\nOriginal query: {}'.format(result["errors"], query))

                if not is_mutation:
                    import inspect
                    caller = inspect.getouterframes(inspect.currentframe(), 2)[1][3]
                    if caller not in list(self.api_costs.keys()):
                        self.api_costs[caller] = 0
                    self.api_costs[caller] += result['data']['rateLimit']['cost']

                return result['data']
            else:
                import json
                raise Exception('Query failed with code {code}:\n{json}'.format(code=request.status_code, json=json.dumps(request.json(), indent=4)))