Move the stuff related to runners to lambda_shared

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]:

tests/ci/ci_runners_metrics_lambda/lambda_shared

@@ -0,0 +1 @@
+../lambda_shared_package/lambda_shared

tests/ci/ci_runners_metrics_lambda/requirements.txt

@@ -1,3 +1,3 @@
-requests<2.30
+../lambda_shared_package
 PyJWT
 cryptography<38

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

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]

tests/ci/terminate_runner_lambda/lambda_shared

@@ -0,0 +1 @@
+../lambda_shared_package/lambda_shared

tests/ci/terminate_runner_lambda/requirements.txt

@@ -1,3 +1,3 @@
-requests<2.30
+../lambda_shared_package
 PyJWT
 cryptography<38