Reduce API calls to SSM client

This commit is contained in:
Mikhail f. Shiryaev 2023-08-23 21:24:33 +02:00
parent bea651e828
commit 0b9e7fa9c5
No known key found for this signature in database
GPG Key ID: 4B02ED204C7D93F4

View File

@ -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