Merge pull request #33990 from ClickHouse/cancel-lambda

Cancel lambda
This commit is contained in:
alesapin 2022-01-26 10:51:21 +03:00 committed by GitHub
commit 677a7f1133
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -5,22 +5,23 @@ import json
import time
import jwt
import requests
import boto3
import requests # type: ignore
import boto3 # type: ignore
NEED_RERUN_OR_CANCELL_WORKFLOWS = {
13241696, # PR
15834118, # Docs
15516108, # ReleaseCI
15797242, # BackportPR
"PullRequestCI",
"Docs",
"DocsRelease",
"BackportPR",
}
# https://docs.github.com/en/rest/reference/actions#cancel-a-workflow-run
#
API_URL = 'https://api.github.com/repos/ClickHouse/ClickHouse'
API_URL = "https://api.github.com/repos/ClickHouse/ClickHouse"
MAX_RETRY = 5
def get_installation_id(jwt_token):
headers = {
"Authorization": f"Bearer {jwt_token}",
@ -29,29 +30,33 @@ def get_installation_id(jwt_token):
response = requests.get("https://api.github.com/app/installations", headers=headers)
response.raise_for_status()
data = response.json()
return data[0]['id']
return data[0]["id"]
def get_access_token(jwt_token, installation_id):
headers = {
"Authorization": f"Bearer {jwt_token}",
"Accept": "application/vnd.github.v3+json",
}
response = requests.post(f"https://api.github.com/app/installations/{installation_id}/access_tokens", headers=headers)
response = requests.post(
f"https://api.github.com/app/installations/{installation_id}/access_tokens",
headers=headers,
)
response.raise_for_status()
data = response.json()
return data['token']
return data["token"]
def get_key_and_app_from_aws():
secret_name = "clickhouse_github_secret_key"
session = boto3.session.Session()
client = session.client(
service_name='secretsmanager',
service_name="secretsmanager",
)
get_secret_value_response = client.get_secret_value(
SecretId=secret_name
)
data = json.loads(get_secret_value_response['SecretString'])
return data['clickhouse-app-key'], int(data['clickhouse-app-id'])
get_secret_value_response = client.get_secret_value(SecretId=secret_name)
data = json.loads(get_secret_value_response["SecretString"])
return data["clickhouse-app-key"], int(data["clickhouse-app-id"])
def get_token_from_aws():
private_key, app_id = get_key_and_app_from_aws()
@ -65,6 +70,7 @@ def get_token_from_aws():
installation_id = get_installation_id(encoded_jwt)
return get_access_token(encoded_jwt, installation_id)
def _exec_get_with_retry(url):
for i in range(MAX_RETRY):
try:
@ -78,20 +84,25 @@ def _exec_get_with_retry(url):
raise Exception("Cannot execute GET request with retries")
WorkflowDescription = namedtuple('WorkflowDescription',
['run_id', 'status', 'rerun_url', 'cancel_url'])
WorkflowDescription = namedtuple(
"WorkflowDescription", ["run_id", "status", "rerun_url", "cancel_url"]
)
def get_workflows_description_for_pull_request(pull_request_event):
head_branch = pull_request_event['head']['ref']
print("PR", pull_request_event['number'], "has head ref", head_branch)
head_branch = pull_request_event["head"]["ref"]
print("PR", pull_request_event["number"], "has head ref", head_branch)
workflows_data = []
workflows = _exec_get_with_retry(API_URL + f"/actions/runs?branch={head_branch}&event=pull_request&page=1")
workflows_data += workflows['workflow_runs']
workflows = _exec_get_with_retry(
API_URL + f"/actions/runs?branch={head_branch}&event=pull_request&page=1"
)
workflows_data += workflows["workflow_runs"]
i = 2
while len(workflows['workflow_runs']) > 0:
workflows = _exec_get_with_retry(API_URL + f"/actions/runs?branch={head_branch}&event=pull_request&page={i}")
workflows_data += workflows['workflow_runs']
while len(workflows["workflow_runs"]) > 0:
workflows = _exec_get_with_retry(
API_URL + f"/actions/runs?branch={head_branch}&event=pull_request&page={i}"
)
workflows_data += workflows["workflow_runs"]
i += 1
if i > 30:
print("Too many workflows found")
@ -99,29 +110,37 @@ def get_workflows_description_for_pull_request(pull_request_event):
workflow_descriptions = []
for workflow in workflows_data:
# unfortunately we cannot filter workflows from forks in request to API so doing it manually
if (workflow['head_repository']['full_name'] == pull_request_event['head']['repo']['full_name']
and workflow['workflow_id'] in NEED_RERUN_OR_CANCELL_WORKFLOWS):
workflow_descriptions.append(WorkflowDescription(
run_id=workflow['id'],
status=workflow['status'],
rerun_url=workflow['rerun_url'],
cancel_url=workflow['cancel_url']))
# unfortunately we cannot filter workflows from forks in request to API
# so doing it manually
if (
workflow["head_repository"]["full_name"]
== pull_request_event["head"]["repo"]["full_name"]
and workflow["name"] in NEED_RERUN_OR_CANCELL_WORKFLOWS
):
workflow_descriptions.append(
WorkflowDescription(
run_id=workflow["id"],
status=workflow["status"],
rerun_url=workflow["rerun_url"],
cancel_url=workflow["cancel_url"],
)
)
return workflow_descriptions
def get_workflow_description(workflow_id):
workflow = _exec_get_with_retry(API_URL + f"/actions/runs/{workflow_id}")
return WorkflowDescription(
run_id=workflow['id'],
status=workflow['status'],
rerun_url=workflow['rerun_url'],
cancel_url=workflow['cancel_url'])
run_id=workflow["id"],
status=workflow["status"],
rerun_url=workflow["rerun_url"],
cancel_url=workflow["cancel_url"],
)
def _exec_post_with_retry(url, token):
headers = {
"Authorization": f"token {token}"
}
headers = {"Authorization": f"token {token}"}
for i in range(MAX_RETRY):
try:
response = requests.post(url, headers=headers)
@ -133,32 +152,34 @@ def _exec_post_with_retry(url, token):
raise Exception("Cannot execute POST request with retry")
def exec_workflow_url(urls_to_cancel, token):
for url in urls_to_cancel:
print("Post for workflow workflow using url", url)
_exec_post_with_retry(url, token)
print("Workflow post finished")
def main(event):
token = get_token_from_aws()
event_data = json.loads(event['body'])
event_data = json.loads(event["body"])
print("Got event for PR", event_data['number'])
action = event_data['action']
print("Got action", event_data['action'])
pull_request = event_data['pull_request']
labels = { l['name'] for l in pull_request['labels'] }
print("Got event for PR", event_data["number"])
action = event_data["action"]
print("Got action", event_data["action"])
pull_request = event_data["pull_request"]
labels = {label["name"] for label in pull_request["labels"]}
print("PR has labels", labels)
if action == 'closed' or 'do not test' in labels:
if action == "closed" or "do not test" in labels:
print("PR merged/closed or manually labeled 'do not test' will kill workflows")
workflow_descriptions = get_workflows_description_for_pull_request(pull_request)
urls_to_cancel = []
for workflow_description in workflow_descriptions:
if workflow_description.status != 'completed':
if workflow_description.status != "completed":
urls_to_cancel.append(workflow_description.cancel_url)
print(f"Found {len(urls_to_cancel)} workflows to cancel")
exec_workflow_url(urls_to_cancel, token)
elif action == 'labeled' and 'can be tested' in labels:
elif action == "labeled" and "can be tested" in labels:
print("PR marked with can be tested label, rerun workflow")
workflow_descriptions = get_workflows_description_for_pull_request(pull_request)
if not workflow_descriptions:
@ -168,7 +189,7 @@ def main(event):
sorted_workflows = list(sorted(workflow_descriptions, key=lambda x: x.run_id))
most_recent_workflow = sorted_workflows[-1]
print("Latest workflow", most_recent_workflow)
if most_recent_workflow.status != 'completed':
if most_recent_workflow.status != "completed":
print("Latest workflow is not completed, cancelling")
exec_workflow_url([most_recent_workflow.cancel_url], token)
print("Cancelled")
@ -176,7 +197,7 @@ def main(event):
for _ in range(30):
latest_workflow_desc = get_workflow_description(most_recent_workflow.run_id)
print("Checking latest workflow", latest_workflow_desc)
if latest_workflow_desc.status in ('completed', 'cancelled'):
if latest_workflow_desc.status in ("completed", "cancelled"):
print("Finally latest workflow done, going to rerun")
exec_workflow_url([most_recent_workflow.rerun_url], token)
print("Rerun finished, exiting")
@ -187,5 +208,6 @@ def main(event):
else:
print("Nothing to do")
def handler(event, _):
main(event)