diff --git a/tests/ci/ci_runners_metrics_lambda/app.py b/tests/ci/ci_runners_metrics_lambda/app.py index 341e1b674ec..a12143752a1 100644 --- a/tests/ci/ci_runners_metrics_lambda/app.py +++ b/tests/ci/ci_runners_metrics_lambda/app.py @@ -10,7 +10,6 @@ import argparse import sys import json import time -from collections import namedtuple from datetime import datetime from typing import Dict, List, Tuple @@ -19,21 +18,14 @@ import requests # type: ignore import boto3 # type: ignore from botocore.exceptions import ClientError # type: ignore -UNIVERSAL_LABEL = "universal" -RUNNER_TYPE_LABELS = [ - "builder", - "func-tester", - "func-tester-aarch64", - "fuzzer-unit-tester", - "stress-tester", - "style-checker", - "style-checker-aarch64", -] - -RunnerDescription = namedtuple( - "RunnerDescription", ["id", "name", "tags", "offline", "busy"] +from lambda_shared import ( + RUNNER_TYPE_LABELS, + RunnerDescription, + RunnerDescriptions, + list_runners, ) -RunnerDescriptions = List[RunnerDescription] + +UNIVERSAL_LABEL = "universal" def get_dead_runners_in_ec2(runners: RunnerDescriptions) -> RunnerDescriptions: @@ -193,52 +185,6 @@ def get_access_token(jwt_token: str, installation_id: int) -> str: return data["token"] # type: ignore -def list_runners(access_token: str) -> RunnerDescriptions: - headers = { - "Authorization": f"token {access_token}", - "Accept": "application/vnd.github.v3+json", - } - per_page = 100 - response = requests.get( - f"https://api.github.com/orgs/ClickHouse/actions/runners?per_page={per_page}", - headers=headers, - ) - response.raise_for_status() - data = response.json() - total_runners = data["total_count"] - print("Expected total runners", total_runners) - runners = data["runners"] - - # round to 0 for 0, 1 for 1..100, but to 2 for 101..200 - total_pages = (total_runners - 1) // per_page + 1 - - print("Total pages", total_pages) - for i in range(2, total_pages + 1): - response = requests.get( - "https://api.github.com/orgs/ClickHouse/actions/runners" - f"?page={i}&per_page={per_page}", - headers=headers, - ) - response.raise_for_status() - data = response.json() - runners += data["runners"] - - print("Total runners", len(runners)) - result = [] - for runner in runners: - tags = [tag["name"] for tag in runner["labels"]] - desc = RunnerDescription( - id=runner["id"], - name=runner["name"], - tags=tags, - offline=runner["status"] == "offline", - busy=runner["busy"], - ) - result.append(desc) - - return result - - def group_runners_by_tag( listed_runners: RunnerDescriptions, ) -> Dict[str, RunnerDescriptions]: diff --git a/tests/ci/ci_runners_metrics_lambda/lambda_shared b/tests/ci/ci_runners_metrics_lambda/lambda_shared new file mode 120000 index 00000000000..ba86e090f6c --- /dev/null +++ b/tests/ci/ci_runners_metrics_lambda/lambda_shared @@ -0,0 +1 @@ +../lambda_shared_package/lambda_shared \ No newline at end of file diff --git a/tests/ci/ci_runners_metrics_lambda/requirements.txt b/tests/ci/ci_runners_metrics_lambda/requirements.txt index 98be09ab232..e99dee1743c 100644 --- a/tests/ci/ci_runners_metrics_lambda/requirements.txt +++ b/tests/ci/ci_runners_metrics_lambda/requirements.txt @@ -1,3 +1,3 @@ -requests<2.30 +../lambda_shared_package PyJWT cryptography<38 diff --git a/tests/ci/lambda_shared_package/lambda_shared/__init__.py b/tests/ci/lambda_shared_package/lambda_shared/__init__.py index c5ae4df9e17..fe52f98d5f6 100644 --- a/tests/ci/lambda_shared_package/lambda_shared/__init__.py +++ b/tests/ci/lambda_shared_package/lambda_shared/__init__.py @@ -4,7 +4,8 @@ It exists as __init__.py and lambda_shared/__init__.py to work both in local and import json import logging import time -from typing import List, Optional +from collections import namedtuple +from typing import Any, List, Optional import boto3 # type: ignore import requests # type: ignore @@ -21,10 +22,14 @@ RUNNER_TYPE_LABELS = [ ### VENDORING -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(Name=name, WithDecryption=decrypt)[ # type: ignore + "Parameter" + ]["Value"] class CHException(Exception): @@ -65,10 +70,64 @@ class ClickHouseHelper: raise CHException("Cannot fetch data from clickhouse") - def select_json_each_row(self, db: str, query: str) -> List[dict]: # type: ignore + def select_json_each_row(self, db: str, query: str) -> List[dict]: text = self._select_and_get_json_each_row(db, query) result = [] for line in text.split("\n"): if line: result.append(json.loads(line)) return result + + +### Runners + +RunnerDescription = namedtuple( + "RunnerDescription", ["id", "name", "tags", "offline", "busy"] +) +RunnerDescriptions = List[RunnerDescription] + + +def list_runners(access_token: str) -> RunnerDescriptions: + headers = { + "Authorization": f"token {access_token}", + "Accept": "application/vnd.github.v3+json", + } + per_page = 100 + response = requests.get( + f"https://api.github.com/orgs/ClickHouse/actions/runners?per_page={per_page}", + headers=headers, + ) + response.raise_for_status() + data = response.json() + total_runners = data["total_count"] + print("Expected total runners", total_runners) + runners = data["runners"] + + # round to 0 for 0, 1 for 1..100, but to 2 for 101..200 + total_pages = (total_runners - 1) // per_page + 1 + + print("Total pages", total_pages) + for i in range(2, total_pages + 1): + response = requests.get( + "https://api.github.com/orgs/ClickHouse/actions/runners" + f"?page={i}&per_page={per_page}", + headers=headers, + ) + response.raise_for_status() + data = response.json() + runners += data["runners"] + + print("Total runners", len(runners)) + result = [] + for runner in runners: + tags = [tag["name"] for tag in runner["labels"]] + desc = RunnerDescription( + id=runner["id"], + name=runner["name"], + tags=tags, + offline=runner["status"] == "offline", + busy=runner["busy"], + ) + result.append(desc) + + return result diff --git a/tests/ci/terminate_runner_lambda/app.py b/tests/ci/terminate_runner_lambda/app.py index e0164bc58c0..5799a498d5a 100644 --- a/tests/ci/terminate_runner_lambda/app.py +++ b/tests/ci/terminate_runner_lambda/app.py @@ -4,7 +4,6 @@ import argparse import json import sys import time -from collections import namedtuple from dataclasses import dataclass from typing import Any, Dict, List, Tuple @@ -12,6 +11,8 @@ import boto3 # type: ignore import requests # type: ignore import jwt +from lambda_shared import RunnerDescriptions, list_runners + def get_key_and_app_from_aws() -> Tuple[str, int]: secret_name = "clickhouse_github_secret_key" @@ -118,58 +119,6 @@ def get_cached_instances() -> dict: return cached_instances.value -RunnerDescription = namedtuple( - "RunnerDescription", ["id", "name", "tags", "offline", "busy"] -) -RunnerDescriptions = List[RunnerDescription] - - -def list_runners(access_token: str) -> RunnerDescriptions: - headers = { - "Authorization": f"token {access_token}", - "Accept": "application/vnd.github.v3+json", - } - per_page = 100 - response = requests.get( - f"https://api.github.com/orgs/ClickHouse/actions/runners?per_page={per_page}", - headers=headers, - ) - response.raise_for_status() - data = response.json() - total_runners = data["total_count"] - print("Expected total runners", total_runners) - runners = data["runners"] - - # round to 0 for 0, 1 for 1..100, but to 2 for 101..200 - total_pages = (total_runners - 1) // per_page + 1 - - print("Total pages", total_pages) - for i in range(2, total_pages + 1): - response = requests.get( - "https://api.github.com/orgs/ClickHouse/actions/runners" - f"?page={i}&per_page={per_page}", - headers=headers, - ) - response.raise_for_status() - data = response.json() - runners += data["runners"] - - print("Total runners", len(runners)) - result = [] - for runner in runners: - tags = [tag["name"] for tag in runner["labels"]] - desc = RunnerDescription( - id=runner["id"], - name=runner["name"], - tags=tags, - offline=runner["status"] == "offline", - busy=runner["busy"], - ) - result.append(desc) - - return result - - def how_many_instances_to_kill(event_data: dict) -> Dict[str, int]: data_array = event_data["CapacityToTerminate"] to_kill_by_zone = {} # type: Dict[str, int] diff --git a/tests/ci/terminate_runner_lambda/lambda_shared b/tests/ci/terminate_runner_lambda/lambda_shared new file mode 120000 index 00000000000..ba86e090f6c --- /dev/null +++ b/tests/ci/terminate_runner_lambda/lambda_shared @@ -0,0 +1 @@ +../lambda_shared_package/lambda_shared \ No newline at end of file diff --git a/tests/ci/terminate_runner_lambda/requirements.txt b/tests/ci/terminate_runner_lambda/requirements.txt index 98be09ab232..e99dee1743c 100644 --- a/tests/ci/terminate_runner_lambda/requirements.txt +++ b/tests/ci/terminate_runner_lambda/requirements.txt @@ -1,3 +1,3 @@ -requests<2.30 +../lambda_shared_package PyJWT cryptography<38