ClickHouse/tests/ci/cherry_pick_utils/query.py

423 lines
14 KiB
Python
Raw Normal View History

2021-11-08 14:30:27 +00:00
# -*- 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)))