ClickHouse/tests/ci/get_robot_token.py

111 lines
3.1 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
import logging
import random
from dataclasses import dataclass
2023-10-04 11:53:55 +00:00
from typing import Any, Dict, List, Optional, Union
import boto3 # type: ignore
2022-11-10 16:11:23 +00:00
from github import Github
from github.AuthenticatedUser import AuthenticatedUser
2023-12-30 19:26:24 +00:00
from github.GithubException import BadCredentialsException
2023-10-04 11:53:55 +00:00
from github.NamedUser import NamedUser
2022-11-10 16:11:23 +00:00
@dataclass
class Token:
2023-10-04 11:53:55 +00:00
user: Union[AuthenticatedUser, NamedUser]
2022-11-10 16:11:23 +00:00
value: str
rest: int
SAFE_REQUESTS_LIMIT = 1000
2023-08-23 19:24:33 +00:00
def get_parameter_from_ssm(
name: str, decrypt: bool = True, client: Optional[Any] = None
) -> str:
2021-10-20 11:48:27 +00:00
if not client:
client = boto3.client("ssm", region_name="us-east-1")
2023-08-23 19:24:33 +00:00
return client.get_parameter( # type:ignore
Name=name, WithDecryption=decrypt
)[
"Parameter"
]["Value"]
def get_parameters_from_ssm(
names: List[str], decrypt: bool = True, client: Optional[Any] = None
) -> Dict[str, str]:
if not client:
client = boto3.client("ssm", region_name="us-east-1")
names = list(set(names))
results = {} # type: Dict[str,str]
i = 0
while (i) * 10 < len(names):
# the get_parameters returns up to 10 values, so the call is split by 10
results.update(
**{
p["Name"]: p["Value"]
for p in client.get_parameters(
Names=names[i * 10 : (i + 1) * 10], WithDecryption=decrypt
)["Parameters"]
}
)
i += 1
return results
2021-10-20 11:48:27 +00:00
2023-04-28 14:01:53 +00:00
ROBOT_TOKEN = None # type: Optional[Token]
def get_best_robot_token(tokens_path: str = "/github-tokens") -> str:
2023-04-28 14:01:53 +00:00
global ROBOT_TOKEN
if ROBOT_TOKEN is not None:
return ROBOT_TOKEN.value
client = boto3.client("ssm", region_name="us-east-1")
tokens = {
p["Name"]: p["Value"]
for p in client.get_parameters_by_path(Path=tokens_path, WithDecryption=True)[
"Parameters"
]
}
assert tokens
token_items = list(tokens.items())
random.shuffle(token_items)
best_token: Optional[Token] = None
for name, value in token_items:
gh = Github(value, per_page=100)
2023-12-30 19:26:24 +00:00
try:
# Do not spend additional request to API by accessing user.login unless
2023-12-30 19:26:24 +00:00
# the token is chosen by the remaining requests number
user = gh.get_user()
rest, _ = gh.rate_limiting
except BadCredentialsException:
logging.error(
"The token %(name)s has expired, please update it\n"
"::error::Token %(name)s has expired, it must be updated",
{"name": name},
)
continue
2022-11-22 12:37:08 +00:00
logging.info("Get token with %s remaining requests", rest)
if best_token is None:
best_token = Token(user, value, rest)
elif best_token.rest < rest:
best_token = Token(user, value, rest)
if best_token.rest > SAFE_REQUESTS_LIMIT:
break
assert best_token
ROBOT_TOKEN = best_token
logging.info(
2023-04-28 14:01:53 +00:00
"User %s with %s remaining requests is used",
ROBOT_TOKEN.user.login,
ROBOT_TOKEN.rest,
)
2023-04-28 14:01:53 +00:00
return ROBOT_TOKEN.value