Merge pull request #53762 from ClickHouse/reduce-get-parameter-calls

Reduce API calls to SSM client

tests/ci/get_robot_token.py

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