mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-11-10 01:25:21 +00:00
Reduce API calls to SSM client
This commit is contained in:
parent
bea651e828
commit
0b9e7fa9c5
@ -1,7 +1,7 @@
|
||||
#!/usr/bin/env python3
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import boto3 # type: ignore
|
||||
from github import Github
|
||||
@ -15,10 +15,40 @@ class Token:
|
||||
rest: int
|
||||
|
||||
|
||||
def get_parameter_from_ssm(name, decrypt=True, client=None):
|
||||
def get_parameter_from_ssm(
|
||||
name: str, decrypt: bool = True, client: Optional[Any] = None
|
||||
) -> str:
|
||||
if not client:
|
||||
client = boto3.client("ssm", region_name="us-east-1")
|
||||
return client.get_parameter(Name=name, WithDecryption=decrypt)["Parameter"]["Value"]
|
||||
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
|
||||
|
||||
|
||||
ROBOT_TOKEN = None # type: Optional[Token]
|
||||
@ -29,15 +59,15 @@ def get_best_robot_token(token_prefix_env_name="github_robot_token_"):
|
||||
if ROBOT_TOKEN is not None:
|
||||
return ROBOT_TOKEN.value
|
||||
client = boto3.client("ssm", region_name="us-east-1")
|
||||
parameters = client.describe_parameters(
|
||||
parameters_list = client.describe_parameters(
|
||||
ParameterFilters=[
|
||||
{"Key": "Name", "Option": "BeginsWith", "Values": [token_prefix_env_name]}
|
||||
]
|
||||
)["Parameters"]
|
||||
assert parameters
|
||||
assert parameters_list
|
||||
tokens = get_parameters_from_ssm([p["Name"] for p in parameters_list], True, client)
|
||||
|
||||
for token_name in [p["Name"] for p in parameters]:
|
||||
value = get_parameter_from_ssm(token_name, True, client)
|
||||
for value in tokens.values():
|
||||
gh = Github(value, per_page=100)
|
||||
# Do not spend additional request to API by accessin user.login unless
|
||||
# the token is chosen by the remaining requests number
|
||||
|
Loading…
Reference in New Issue
Block a user